TRAK on 5D shapes

[Bas-2k]

Hi,
I have a classification model trained on the shape batch_size x N (=46) x channel x height x width. How can I adapt TRAK to use for that ? I get an error in the featurize function itself. Do I have to modify the in_dims?

/opt/conda/lib/python3.8/site-packages/trak/gradient_computers.py in compute_per_sample_grad(self, batch)
148
149 # map over batch dimensions (hence 0 for each batch dimension, and None for model params)
--> 150 grads = torch.func.vmap(
151 grads_loss,
152 in_dims=(None, None, None, *([0] * len(batch))),

/opt/conda/lib/python3.8/site-packages/torch/_functorch/vmap.py in wrapped(*args, **kwargs)
432
433 # If chunk_size is not specified.
--> 434 return _flat_vmap(
435 func, batch_size, flat_in_dims, flat_args, args_spec, out_dims, randomness, **kwargs
436 )

/opt/conda/lib/python3.8/site-packages/torch/_functorch/vmap.py in fn(*args, **kwargs)
37 def fn(*args, **kwargs):
38 with torch.autograd.graph.disable_saved_tensors_hooks(message):
---> 39 return f(*args, **kwargs)
40 return fn
41

/opt/conda/lib/python3.8/site-packages/torch/_functorch/vmap.py in _flat_vmap(func, batch_size, flat_in_dims, flat_args, args_spec, out_dims, randomness, **kwargs)
617 try:
618 batched_inputs = _create_batched_inputs(flat_in_dims, flat_args, vmap_level, args_spec)
--> 619 batched_outputs = func(*batched_inputs, **kwargs)
620 return _unwrap_batched(batched_outputs, out_dims, vmap_level, batch_size, func)
621 finally:

/opt/conda/lib/python3.8/site-packages/torch/functorch/eager_transforms.py in wrapper(*args, **kwargs)
1378 @wraps(func)
1379 def wrapper(*args, **kwargs):
-> 1380 results = grad_and_value(func, argnums, has_aux=has_aux)(*args, **kwargs)
1381 if has_aux:
1382 grad, (, aux) = results

/opt/conda/lib/python3.8/site-packages/torch/_functorch/vmap.py in fn(*args, **kwargs)
37 def fn(*args, **kwargs):
38 with torch.autograd.graph.disable_saved_tensors_hooks(message):
---> 39 return f(*args, **kwargs)
40 return fn
41

/opt/conda/lib/python3.8/site-packages/torch/functorch/eager_transforms.py in wrapper(*args, **kwargs)
1243 tree_map(partial(_create_differentiable, level=level), diff_args)
1244
-> 1245 output = func(*args, **kwargs)
1246 if has_aux:
1247 if not (isinstance(output, tuple) and len(output) == 2):

/opt/conda/lib/python3.8/site-packages/trak/modelout_functions.py in get_output(model, weights, buffers, image, label)
138 """
139 logits = ch.func.functional_call(model, (weights, buffers), image.unsqueeze(0))
--> 140 bindex = ch.arange(logits.shape[0]).to(logits.device, non_blocking=False)
141 logits_correct = logits[bindex, label.unsqueeze(0)]
142

AttributeError: 'tuple' object has no attribute 'shape'