fixed device issues

inference.py

@@ -9,6 +9,8 @@ import io
 
 from model import Unet, ConditionalUnet, GaussianDiffusion, DiffusionImageAPI
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 def inference1():
   # new image from web page
   image = requests.get("https://picsum.photos/120/80").content
@@ -18,7 +20,7 @@ def inference():
   model = Unet(
     image_channels=3,
   )
-  model.load_state_dict(torch.load("./model_final.pt"))
+  model.load_state_dict(torch.load("./model_final.pt"), map_location=device)
 
   diffusion = GaussianDiffusion(
     model=model,
@@ -28,6 +30,9 @@ def inference():
     image_size=(120, 80),
   )
 
+  model.to(device)
+  diffusion.to(device)
+
   imageAPI = DiffusionImageAPI(diffusion)
 
   images, versions = diffusion.sample(1)