Add torch device selection

main.py

@@ -10,13 +10,17 @@ from PIL import Image
 from models import phc_models
 from utils import utils
 
+device = torch.device('cpu')
+if torch.cuda.is_available():
+    device = torch.device('cuda:0')
+
 BILATERIAL_WEIGHT = 'weights/phresnet18_cbis2views.pt'
 BILATERAL_MODEL = phc_models.PHCResNet18(
     channels=2, n=2, num_classes=1, visualize=True)
 BILATERAL_MODEL.add_top_blocks(num_classes=1)
 BILATERAL_MODEL.load_state_dict(torch.load(
     BILATERIAL_WEIGHT, map_location='cpu'))
-BILATERAL_MODEL = BILATERAL_MODEL.to('cpu')
+BILATERAL_MODEL = BILATERAL_MODEL.to(device)
 BILATERAL_MODEL.eval()
 INPUT_HEIGHT, INPUT_WIDTH = 600, 500
 
@@ -143,6 +147,7 @@ def predict_bilateral(cc_file, mlo_file):
 
     images_t = torch.from_numpy(images)
     images_t = images_t.unsqueeze(0)  # Add batch dimension
+    images_t = images_t.to(device)
 
     out, _, out_refiner = BILATERAL_MODEL(images_t)