Update infer.py

infer.py

@@ -56,7 +56,7 @@ def infer_pipe(pipe, images_batch, task_name, seed, device):
             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(batch_size, 1)
+            task_emb = task_emb.expand(batch_size, -1)
 
             # Run inference
             preds = pipe(