Nicholas Vadivelu

Did anyone find a solution for this?

Will look into it, thanks for bring it to my attention! Would also be open to PRs if you are able.

Hi @sarahelsherif, thanks for raising this issue! Could you also paste in the PyTorch code that gives you `torch.Size ([1, 2048, 28, 28])` for comparison?

Hi @sarahelsherif, I wasn't able to directly use the code that you sent since I do not have `create_feature_extractor`. Instead, see this example of extracting the bacbone in both JAX...

This can definitely be supported, essentially we would need to apply the logic in [`_make_layer`](https://github.com/pytorch/vision/blob/33123bee8554af85989bc7188bf3c7dc1ba5e8c6/torchvision/models/resnet.py#L223-L262) to the [`ResNetBottleneckBlock`](https://github.com/n2cholas/jax-resnet/blob/43f41194123d63ca5e7d388cab8f85ae5c6d0abf/jax_resnet/resnet.py#L101-L127). I won't have the bandwidth to work on this for a few...

Thanks for reporting this! This is an issue on the PyTorch side, as I get the same error with: ```py import torch model = torch.hub.load('pytorch/vision:v0.9.0', 'resnet18', pretrained=True) ``` I opened...

I still get the same error on Colab with: ```py import torch model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=True) ```

The issue is being tracked here: https://github.com/pytorch/vision/issues/4156

@ruotianluo the issue is temporarily fixed via https://github.com/n2cholas/jax-resnet/commit/43f41194123d63ca5e7d388cab8f85ae5c6d0abf

I think the Python saying "We're all consenting adults here" is pretty fitting. In my view, trading convenience for safety is reasonable here because JAX users should know (or will...