Fabric and FFCV?

[richardrl]

Description & Motivation

Does Fabric work with FFCV dataloader?
I just tried it and first, when I used fabric.setup(..., ...), it said fabric only works with Pytorch dataloader.

So I removed it. But then, when running my script, I got this:

File "/data/scratch-oc40/pulkitag/rli14/micromamba/ffcv_oct40/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
  return self._call_impl(*args, **kwargs)
File "/data/scratch-oc40/pulkitag/rli14/micromamba/ffcv_oct40/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
  return forward_call(*args, **kwargs)
File "/data/scratch-oc40/pulkitag/rli14/hamer_diffusion_policy/model/vision/multi_image_obs_encoder.py", line 216, in forward
  feature = self.key_model_map[key](img)
File "/data/scratch-oc40/pulkitag/rli14/micromamba/ffcv_oct40/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
  return self._call_impl(*args, **kwargs)
File "/data/scratch-oc40/pulkitag/rli14/micromamba/ffcv_oct40/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
  return forward_call(*args, **kwargs)
File "/data/scratch-oc40/pulkitag/rli14/micromamba/ffcv_oct40/lib/python3.9/site-packages/torchvision/models/vision_transformer.py", line 291, in forward
  x = self._process_input(x)
File "/data/scratch-oc40/pulkitag/rli14/micromamba/ffcv_oct40/lib/python3.9/site-packages/torchvision/models/vision_transformer.py", line 277, in _process_input
  x = self.conv_proj(x)
File "/data/scratch-oc40/pulkitag/rli14/micromamba/ffcv_oct40/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
  return self._call_impl(*args, **kwargs)
File "/data/scratch-oc40/pulkitag/rli14/micromamba/ffcv_oct40/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
  return forward_call(*args, **kwargs)
File "/data/scratch-oc40/pulkitag/rli14/micromamba/ffcv_oct40/lib/python3.9/site-packages/torch/nn/modules/conv.py", line 460, in forward
  return self._conv_forward(input, self.weight, self.bias)
File "/data/scratch-oc40/pulkitag/rli14/micromamba/ffcv_oct40/lib/python3.9/site-packages/torch/nn/modules/conv.py", line 456, in _conv_forward
  return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

I also set this: torch.set_float32_matmul_precision('medium') # or 'high'

Pitch

I want Fabric and FFCV to work together well.

Alternatives

No response

Additional context

No response

cc @Borda

[lantiga]

Can you please post a minimal repro?