add device = 'cpu'

app.py

@@ -1,9 +1,19 @@
 import torch
 from monai.bundle import ConfigParser
 import gradio as gr
+import json
+
+with open("configs/inference.json") as f:
+    inference_config = json.load(f)
+
+device = torch.device('cpu')
+if torch.cuda.is_available():
+    device = torch.device('cuda:0')
+
+inference_config["device"] = device
 
 parser = ConfigParser() #  load configuration files that specify various parameters for running the MONAI workflow.
-parser.read_config(f="configs/inference.json") # read the config from specified JSON file
+parser.read_config(f=inference_config) # read the config from specified JSON file
 parser.read_meta(f="configs/metadata.json") # read the metadata from specified JSON file
 
 inference = parser.get_parsed_content("inferer")
@@ -11,6 +21,9 @@ network = parser.get_parsed_content("network_def")
 preprocess = parser.get_parsed_content("preprocessing")
 state_dict = torch.load("models/model.pt")
 network.load_state_dict(state_dict, strict=True) #  Loads a model’s parameter dictionary
+network = network.to(device)
+network.eval()
+
 class_names = {
     0: "Other",
     1: "Inflammatory",
@@ -21,6 +34,7 @@ class_names = {
 def classify_image(image_file, label_file):
     data = {"image":image_file, "label":label_file}
     batch = preprocess(data)
+    batch['image'] = batch['image'].to(device)
     network.eval()
     with torch.no_grad():
         pred = inference(batch['image'].unsqueeze(dim=0), network) # expect 4 channels input  (3 RGB, 1 Label mask)