From 3576a7144f2422ad99e109612fe7fe7e69e3c20f Mon Sep 17 00:00:00 2001 From: Marek Otahal Date: Sun, 26 Jan 2020 20:28:20 +0100 Subject: [PATCH] avoid forced setting to torch.cuda.FloatTensor type as explained in https://github.com/dbolya/yolact/issues/256#issuecomment-577689755 Author: @melaanya , Thanks! --- layers/functions/detection.py | 2 +- yolact.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/layers/functions/detection.py b/layers/functions/detection.py index 4e5fd068c..db92ed9fe 100644 --- a/layers/functions/detection.py +++ b/layers/functions/detection.py @@ -207,7 +207,7 @@ def traditional_nms(self, boxes, masks, scores, iou_threshold=0.5, conf_thresh=0 preds = torch.cat([boxes[conf_mask], cls_scores[:, None]], dim=1).cpu().numpy() keep = cnms(preds, iou_threshold) - keep = torch.Tensor(keep, device=boxes.device).long() + keep = torch.Tensor(keep).to(boxes.device).long() idx_lst.append(idx[keep]) cls_lst.append(keep * 0 + _cls) diff --git a/yolact.py b/yolact.py index d83703bb7..8a0ae8c9e 100644 --- a/yolact.py +++ b/yolact.py @@ -245,7 +245,7 @@ def make_priors(self, conv_h, conv_w, device): prior_data += [x, y, w, h] - self.priors = torch.Tensor(prior_data, device=device).view(-1, 4).detach() + self.priors = torch.Tensor(prior_data).view(-1, 4).to(device).detach() self.priors.requires_grad = False self.last_img_size = (cfg._tmp_img_w, cfg._tmp_img_h) self.last_conv_size = (conv_w, conv_h)