Update infer.py

infer.py

@@ -47,10 +47,10 @@ def infer_pipe(pipe, image, task_name, seed, device):
     with torch.no_grad():
         with autocast_ctx:
             # Convert image to tensor
-            img = np.array(image.convert('RGB')).astype(np.float32)
+            img = np.array(image.convert('RGB')).astype(np.float16)
             test_image = torch.tensor(img).permute(2, 0, 1).unsqueeze(0)
             test_image = test_image / 127.5 - 1.0
-            test_image = test_image.to(device).type(torch.float16)
+            test_image = test_image.to(device)
 
             # Create task_emb
             task_emb = torch.tensor([1, 0], device=device, dtype=torch.float16).unsqueeze(0)