Update handler.py

handler.py

@@ -26,6 +26,8 @@ class EndpointHandler:
         model.hf_device_map["upsample"] = model.hf_device_map["swin2sr.embeddings"]
         self.model = Swin2SRForImageSuperResolution.from_pretrained("caidas/swin2SR-classical-sr-x2-64", device_map=model.hf_device_map)
 
+        print(subprocess.run(["nvidia-smi"]))
+    
     def __call__(self, data: Any):
         """
         Args:
@@ -39,6 +41,8 @@ class EndpointHandler:
         inputs = self.processor(image, return_tensors="pt")
         with torch.no_grad():
             outputs = self.model(**inputs)
+
+        print(subprocess.run(["nvidia-smi"]))
         
         output = outputs.reconstruction.data.squeeze().float().cpu().clamp_(0, 1).numpy()
         output = np.moveaxis(output, source=0, destination=-1)