fix : ray base

image_ref/main.py

@@ -22,6 +22,9 @@ def train_duo(model, data_train, optimizer, loss_function, epoch, wandb):
         param.requires_grad = True
 
     for imaer, imana, img_ref, label in data_train:
+        imaer = imaer.float()
+        imana = imana.float()
+        img_ref = img_ref.float()
         label = label.long()
         if torch.cuda.is_available():
             imaer = imaer.cuda()
@@ -55,6 +58,9 @@ def val_duo(model, data_test, loss_function, epoch, wandb):
         param.requires_grad = False
 
     for imaer, imana, img_ref, label in data_test:
+        imaer = imaer.float()
+        imana = imana.float()
+        img_ref = img_ref.float()
         imaer = imaer.transpose(0, 1)
         imana = imana.transpose(0, 1)
         img_ref = img_ref.transpose(0, 1)
@@ -107,7 +113,7 @@ def run_duo(args):
 
     # load model
     model = Classification_model_duo_contrastive(model=args.model, n_class=2)
-    model.double()
+    model.float()
     # load weight
     if args.pretrain_path is not None:
         print('Model weight loaded')
@@ -197,6 +203,9 @@ def make_prediction_duo(model, data, f_name, f_name2):
     soft_max = nn.Softmax(dim=1)
     # iterate over test data
     for imaer, imana, img_ref, label in data:
+        imaer = imaer.float()
+        imana = imana.float()
+        img_ref = img_ref.float()
         imaer = imaer.transpose(0, 1)
         imana = imana.transpose(0, 1)
         img_ref = img_ref.transpose(0, 1)

main_ray.py

@@ -20,7 +20,7 @@ def train_model(config,args):
     # load data
 
 
-    data_train, data_test = load_data_duo(dataset_dir=args.dataset_dir,
+    data_train, data_test = load_data_duo(base_dir=args.dataset_dir,
                                                   batch_size=args.batch_size,
                                                   noise_threshold=config['noise'],
                                                   )