Update single/load.py

single/load.py

@@ -57,10 +57,10 @@ class CombinedNet(nn.Module):
 
 
 # MLSTAC API -----------------------------------------------------------------------
-def example_data(path: pathlib.Path, *args, **kwargs):
+def example_data(path: pathlib.Path, device = "cpu", *args, **kwargs):
     data_f = path / "example_data.safetensor"    
     sample = safetensors.torch.load_file(data_f)
-    return  sample["image"]
+    return  sample["image"].float().unsqueeze(0).to(device)
 
 def trainable_model(path, device: str = "cpu", *args, **kwargs):
     trainable_f = path / "model.safetensor"
@@ -92,6 +92,7 @@ def compiled_model(path, device: str = "cpu", *args, **kwargs):
 
     return cloud_model
 
+
 def display_results(path: pathlib.Path, device: str = "cpu", *args, **kwargs):
     # Load model
     model = compiled_model(path, device, benchmark=True)
@@ -100,13 +101,13 @@ def display_results(path: pathlib.Path, device: str = "cpu", *args, **kwargs):
     probav = example_data(path)
 
     # Run model
-    cloudprobs = model(probav.float().unsqueeze(0).to(device)).squeeze(0).cpu()
+    cloudprobs = model(probav).squeeze().cpu()
 
     #Display results
     fig, ax = plt.subplots(1, 2, figsize=(8, 4))
-    ax[0].imshow(probav[[2, 1, 0]].cpu().detach().numpy().transpose(1, 2, 0))
+    ax[0].imshow(probav[0, [2, 1, 0]].cpu().detach().numpy().transpose(1, 2, 0))
     ax[0].set_title("Input")
-    ax[1].imshow(cloudprobs[0].cpu().detach().numpy(), cmap="gray")
+    ax[1].imshow(cloudprobs.cpu().detach().numpy(), cmap="gray")
     ax[1].set_title("Output")    
     for a in ax:
         a.axis("off")