TypeError: no implementation found for 'torch.nn.functional.cross_entropy' on types that implement __torch_function__: [<class 'fastai.torch_core.TensorImage'>, <class 'fastai.torch_core.TensorCategory'>]

I am using fastai version 2.5.4, fastcore version 1.3.28 and torch version 1.10.0+cu111
When trying a simple training loop with a SE_Net (based off a ResNet) I keep getting the error
TypeError: no implementation found for ‘torch.nn.functional.cross_entropy’ on types that implement torch_function: [<class ‘fastai.torch_core.TensorImage’>, <class ‘fastai.torch_core.TensorCategory’>]

the error seems to be from:
/usr/local/lib/python3.7/dist-packages/torch/overrides.py in handle_torch_function(public_api, relevant_args, *args, **kwargs)

I have a simple cross_entropy loss defined in a get_learner function as so:

def get_learner(m):

return Learner(dls, m, loss_func=nn.CrossEntropyLoss(), metrics=accuracy).to_fp16()

I am able to load the dataloaders fine and show them using IMAGENETTE_160

Any idea of what I could be doing wrong?? Thanks in advance for any tips or suggestions

Could you switch to CrossEntropyLossFlat instead? Loss Functions | fastai

Thanks Manikandan,
Will try that.

No problem @Srinivas . I noticed the same issue was raised in this post as well. TypeError: no implementation found for 'torch.nn.functional.cross_entropy' - #4 by muellerzr