Update infer.py

infer.py

@@ -47,14 +47,16 @@ def infer_pipe(pipe, images_batch, task_name, seed, device):
     with torch.no_grad():
         with autocast_ctx:
             # Convert list of images to tensor
-            images = [np.array(img.convert('RGB')).astype(np.float16) for img in images_batch]
+            images = [np.array(img.convert('RGB')).astype(np.float32) for img in images_batch]
             test_images = torch.stack([torch.tensor(img).permute(2, 0, 1) for img in images])
             test_images = test_images / 127.5 - 1.0
-            test_images = test_images.to(device)
+            test_images = test_images.to(device).type(torch.float16)
 
-            task_emb = torch.tensor([1, 0]).float().unsqueeze(0).to(device)
+            # Ensure task_emb matches expected dimensions
+            batch_size = test_images.shape[0]
+            task_emb = torch.tensor([1, 0], device=device, dtype=torch.float16).unsqueeze(0)
             task_emb = torch.cat([torch.sin(task_emb), torch.cos(task_emb)], dim=-1)
-            task_emb = task_emb.repeat(len(test_images), 1)
+            task_emb = task_emb.repeat(batch_size, 1)
 
             # Run inference
             preds = pipe(