Hotfix dtype casting issue

Author: synml

demo.py

@@ -38,10 +38,10 @@
     os.makedirs(result_dir, exist_ok=True)
     os.makedirs(groundtruth_dir, exist_ok=True)
     for images, masks in tqdm.tqdm(testloader, desc='Demo'):
-        images = images.to(device)
-
         # mask에 255를 곱하여 0~1 사이의 값을 0~255 값으로 변경 + 채널 차원 제거
-        masks.mul_(255).squeeze_(dim=1).type(torch.LongTensor)
+        masks.mul_(255).squeeze_(dim=1)
+
+        images, masks = images.to(device), masks.type(torch.LongTensor)
 
         # 예측
         with torch.no_grad():

eval.py

@@ -50,11 +50,11 @@ def evaluate(model, testloader, criterion, num_classes: int, device):
     val_loss = 0
     inference_time = 0
     for images, masks in tqdm.tqdm(testloader, desc='Eval', leave=False):
-        images, masks = images.to(device), masks.to(device, dtype=torch.int64)
-
         # mask에 255를 곱하여 0~1 사이의 값을 0~255 값으로 변경 + 채널 차원 제거
         masks.mul_(255).squeeze_(dim=1)
 
+        images, masks = images.to(device), masks.to(device, dtype=torch.int64)
+
         # 예측
         with torch.no_grad():
             start_time = time.time()

train.py

@@ -38,11 +38,11 @@
         model.train()
 
         for batch_idx, (images, masks) in enumerate(tqdm.tqdm(trainloader, desc='Train', leave=False)):
-            images, masks = images.to(device), masks.to(device, dtype=torch.int64)
-
             # mask에 255를 곱하여 0~1 사이의 값을 0~255 값으로 변경 + 채널 차원 제거
             masks.mul_(255).squeeze_(dim=1)
 
+            images, masks = images.to(device), masks.to(device, dtype=torch.int64)
+
             # 순전파 + 역전파 + 최적화
             optimizer.zero_grad()
             masks_pred = model(images)