fix : float double error

image_ref/main_ray.py

@@ -38,7 +38,7 @@ def train_model(config,args):
     model = Classification_model_duo_contrastive(model=args.model, n_class=2)
 
     # move parameters to GPU
-    model.double()
+    model.float()
     device = "cpu"
     if torch.cuda.is_available():
         device = "cuda:0"
@@ -53,6 +53,7 @@ def train_model(config,args):
     # init training
     n_class = len(data_train.dataset.classes)
     weight = torch.Tensor([1/n_class,1-1/n_class])
+    weight.float()
     if torch.cuda.is_available():
         weight = weight.cuda()
     print('weight',weight.device)
@@ -79,6 +80,9 @@ def train_model(config,args):
             param.requires_grad = True
 
         for imaer, imana, img_ref, label in data_train:
+            imaer.float()
+            imana.float()
+            img_ref.float()
             label = label.long()
             if torch.cuda.is_available():
                 imaer = imaer.cuda()
@@ -108,6 +112,9 @@ def train_model(config,args):
             param.requires_grad = False
 
         for imaer, imana, img_ref, label in data_val_batch:
+            imaer.float()
+            imana.float()
+            img_ref.float()
             imaer = imaer.transpose(0, 1)
             imana = imana.transpose(0, 1)
             img_ref = img_ref.transpose(0, 1)