diff --git a/dataset/dataset.py b/dataset/dataset.py index 0ffebbb4dec481e5d0e04f7bc7f27939e2f4cd15..32e96f4396eccfe938b5fc561bbab19e93ba6747 100644 --- a/dataset/dataset.py +++ b/dataset/dataset.py @@ -260,6 +260,8 @@ class ImageFolderDuoPatch(data.Dataset): impathAER, impathANA, target = self.imlist[index] imgAER = self.loader(impathAER) imgANA = self.loader(impathANA) + imgANA = imgANA.to(torch.float32) + imgAER = imgAER.to(torch.float32) shape_ana = imgANA.shape shape_aer = imgAER.shape if self.transform is not None: diff --git a/main.py b/main.py index 42a89ef28c27bd07d923f661c6c1fb44ab770977..b4832c5f05f3afd29b5023e9fc8722ae23a35e9b 100644 --- a/main.py +++ b/main.py @@ -209,7 +209,7 @@ def run_duo(args): else : model = Classification_model_duo_pretrained(model = args.model, n_class=len(data_train.dataset.classes)) - model = model.double() + model = model.to(torch.float32) #load weight if args.pretrain_path is not None : load_model(model,args.pretrain_path)