Updated Cuda again

Model_Class.py

@@ -67,25 +67,25 @@ model.eval()
 
 def load_and_classify_image(image_path, device):
 
-    model = model.to(device)
+    gpu_model = model.to(device)
     image = val_transforms_416x628(image_path)
     image = image.unsqueeze(0).to(device)
 
     with torch.no_grad():
-        prediction = model(image)
+        prediction = gpu_model(image)
         prediction = torch.nn.functional.softmax(prediction, dim=1).squeeze(0)
         return prediction.to('cpu'), image.to('cpu')
     
 
 def make_GradCAM(image, device):
 
-    model = model.to(device)
+    gpu_model = model.to(device)
     image = image.to(device)
     model.eval()
-    target_layers = [model.model.layer4[-1]]
+    target_layers = [gpu_model.gpu_model.layer4[-1]]
 
     arr = image.numpy().squeeze()
-    cam = GradCAM(model=model, target_layers=target_layers)
+    cam = GradCAM(model=gpu_model, target_layers=target_layers)
     targets = None
     grayscale_cam = cam(
         input_tensor=image,

Model_Seg.py

@@ -73,14 +73,13 @@ post_transforms = Compose([
 
 def load_and_segment_image(input_image_path, device):
         
-    model = model.to(device)
     image_tensor = pre_transforms(input_image_path)
     image_tensor = image_tensor.unsqueeze(0).to(device)
 
     # Inference using SlidingWindowInferer
     inferer = SlidingWindowInferer(roi_size=(512, 512), sw_batch_size=16, overlap=0.75)
     with torch.no_grad():
-        outputs = inferer(image_tensor, model)
+        outputs = inferer(image_tensor, model.to(device))
 
 
     outputs = outputs.squeeze(0)