Error in GraspAffordanceHarness

[dtch1997]

Firstly, thanks very much for open-sourcing your package code and evaluation harness!
I'm having trouble running the example and wanted to check if I'm doing something wrong.

I downloaded ARC Grasping dataset and ran the code snippet provided in examples/evaluate_grasp.py.

However, I got this error message:

EinopsError:  Error while processing rearrange-reduction pattern "bsz (h w) d -> bsz d h w".
 Input tensor shape: torch.Size([169, 384]). Additional info: {'h': 1, 'w': 1}.
 Expected 3 dimensions, got 2

System details: Ubuntu 22.04, Python 3.8.16
Versions:

torch = "1.13.1"
torchvision = "0.14.1"
voltron-robotics = "1.0"
voltron-evaluation = {git = "https://github.com/siddk/voltron-evaluation.git"}

You can see that it starts running correctly:

$ python examples/09_intro_voltron_eval.py 
06/12 [11:10:59] INFO     | >> Initializing GraspAffordanceHarness                      harness.py:49
                 INFO     | >> Invoking GraspAffordanceHarness.fit()                    harness.py:65
                 INFO     | >> Starting Data Processing for Fold 1 / 5                  harness.py:74
                 INFO     | >> Instantiating Adapter Model and Callbacks                harness.py:77
                 INFO     | >> Training...                                              harness.py:87
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
                 INFO     | >> Compiling Grasp Train Dataset                     preprocessing.py:133
06/12 [11:11:07] INFO     | >> Compiling Grasp Validation Dataset                preprocessing.py:143
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]                                                            

  | Name        | Type       | Params
-------------------------------------------
0 | backbone    | VCond      | 99.2 M
1 | extractor   | MAPBlock   | 2.5 M 
2 | pup_decoder | PuPDecoder | 540 K 
-------------------------------------------
3.1 M     Trainable params
99.2 M    Non-trainable params
102 M     Total params
409.164   Total estimated model params size (MB)
Sanity Checking DataLoader 0:   0%|                                            | 0/1 [00:00<?, ?it/s]

This is the part of voltron-evaluation where the error occurs

voltron_evaluation/grasping/adapter.py:94 in forward  
│    91 │   │   extracted = self.extractor(patches)                                                │
│    92 │   │   if len(extracted.shape) == 1:                                                      │
│    93 │   │   │   extracted = extracted.unsqueeze(1)                                             │
│ ❱  94 │   │   grid = rearrange(extracted, "bsz (h w) d -> bsz d h w", h=self.grid_size, w=self   │
│    95 │   │                                                                                      │
│    96 │   │   # Get Segmentation logits                                                          │
│    97 │   │   return self.pup_decoder(grid)