Update model/DCAMA.py

model/DCAMA.py

@@ -146,7 +146,7 @@ class DCAMA(nn.Module):
             # n_support_feats = [n_support_feats[arg] for arg in args]
             n_simis = n_simi[args].max()
         else:
-            n_simis = torch.tensor(0.).float().cuda()
+            n_simis = torch.tensor(0.).float()
         return logit_mask, n_simis
     
     def extract_feats(self, img):
@@ -453,11 +453,11 @@ class DCAMA(nn.Module):
                 n_support_feats4 = []
                 for idx in selected:
                     n_support_feats4.append(n_support_feats[idx])
-                support_masks = support_masks[:, torch.tensor(selected).long().cuda(), :, :]
+                support_masks = support_masks[:, torch.tensor(selected).long(), :, :]
                 n_support_feats = n_support_feats4
                 simi_map = None
             else:
-                n_simis = torch.tensor(0.).float().cuda()
+                n_simis = torch.tensor(0.).float()
                 simi_map = None
             
             logit_mask = self.model(query_feats, n_support_feats, support_masks.clone(), nshot)