Create Learner using Pytorch model with replaced head --> Get error "RuntimeError: CUDA error: device-side assert triggered"

import torchvision.models as models

model=models.mobilenet_v2(pretrained=True, progress=True)
for param in model.parameters():
    param.requires_grad = False
new_head = nn.Sequential(nn.Dropout(p=0.2,inplace=False),
                         nn.Linear(in_features=1280,out_features=1000,bias=True),
                         nn.ReLU(),
                         nn.Linear(in_features=1000,out_features=1,bias=True),
                         nn.Sigmoid())
model.classifier = new_head


learn = Learner(hot_databunch, model, metrics=[accuracy])
learn.fit(1)

Get error “RuntimeError: CUDA error: device-side assert triggered”

while

mport torchvision.models as models

model=models.mobilenet_v2(pretrained=True, progress=True)
for param in model.parameters():
    param.requires_grad = False

learn = Learner(hot_databunch, model, metrics=[accuracy])
learn.fit(1)

works fine.

What am I doing wrong here ?

Never mind. Already figured out what is the problem.

Databunch has CrossEntrophy loss and the dimension of my label was wrong.

learn = Learner(hot_databunch, model.cuda(), metrics=[accuracy])