From 7f0c0baf605c70e6dd243bdfddf77bdb55d47860 Mon Sep 17 00:00:00 2001 From: sosokker Date: Mon, 13 May 2024 14:42:47 +0700 Subject: [PATCH] Update scaled anchore dtype and fix cuda as device --- StreamServer/src/analytic/action/Detection/Models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/StreamServer/src/analytic/action/Detection/Models.py b/StreamServer/src/analytic/action/Detection/Models.py index d12d5df..8de38a1 100644 --- a/StreamServer/src/analytic/action/Detection/Models.py +++ b/StreamServer/src/analytic/action/Detection/Models.py @@ -118,7 +118,7 @@ class YOLOLayer(nn.Module): # Calculate offsets for each grid self.grid_x = torch.arange(g).repeat(g, 1).view([1, 1, g, g]).type(FloatTensor) self.grid_y = torch.arange(g).repeat(g, 1).t().view([1, 1, g, g]).type(FloatTensor) - self.scaled_anchors = FloatTensor([(a_w / self.stride, a_h / self.stride) for a_w, a_h in self.anchors]) + self.scaled_anchors = torch.tensor([(a_w / self.stride, a_h / self.stride) for a_w, a_h in self.anchors], dtype=torch.float32, device='cuda') self.anchor_w = self.scaled_anchors[:, 0:1].view((1, self.num_anchors, 1, 1)) self.anchor_h = self.scaled_anchors[:, 1:2].view((1, self.num_anchors, 1, 1))