Remove softmax layer from x3d_l model

I want to remove last softmax layer from pytorch x3d_l model.
This is what I tried

model = torch.hub.load('facebookresearch/pytorchvideo', 'x3d_l', pretrained=True)
model.blocks[5]=nn.Sequential(*list(model.blocks[5].children())[:-2])

Shape of my input is torch.Size([2, 3, 16, 320, 320])
When I pass it to the model, I got this error

RuntimeError: mat1 and mat2 shapes cannot be multiplied (4096x1 and 2048x400)

This last block look like this
enter image description here

After removing last two layer, it look like this
enter image description here

You are currently trying to replace the ResNetBasicHead with a simple nn.Sequential container, which might not work as some calls would be missing from the forward method as seen here. In your particular case I would guess the flattening operation is missing and also be careful about the usage of the “project module” via self.proj since the permutations would also be missing.