Update models/grasp_mods.py

models/grasp_mods.py

@@ -336,6 +336,7 @@ def add_inference_method(self):
         # repeat to batchsize
         grasp_query_pe = grasp_query_pe.repeat(n_queries, 1, 1)
         pixel_masks = torch.cat([batched_input[i]['pixel_mask'] for i in range(len(batched_input))], dim=0)
+        pixel_masks = pixel_masks.repeat(n_queries, 1, 1)
         downsampled_pixel_masks = nn.functional.interpolate(pixel_masks.unsqueeze(1).float(), size=(64, 64),
                                                             mode='nearest').squeeze(1).bool()
         downsampled_pixel_masks = downsampled_pixel_masks.view(batch_size, 64 * 64).contiguous()