RuntimeError: CUDA error: device-side assert triggered CUDA kernel errors might be asynchronously reported at some other API call - image

Question

I tried DeIT classification using pytorch, it's a workable code. When I change my collab account to pro I still get an error in the second epochs:
I execute this instruction :
model_ft = train_model(model, criterion, optimizer, exp_lr_scheduler) # now it is a lot faster
Training model function: :
def train_model(model, criterion, optimizer, scheduler, num_epochs=10):
since = time.time()
best_model_wts = copy.deepcopy(model.state_dict())
best_acc = 0.0
for epoch in range(num_epochs):
print(f'Epoch {epoch}/{num_epochs - 1}')
print("-"*10)
for phase in ['train', 'val']: # We do training and validation phase per epoch
if phase == 'train':
model.train() # model to training mode
else:
model.eval() # model to evaluate
running_loss = 0.0
running_corrects = 0.0
for inputs, labels in tqdm(dataloaders[phase]):
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
with torch.set_grad_enabled(phase == 'train'): # no autograd makes validation go faster
outputs = model(inputs)
_, preds = torch.max(outputs, 1) # used for accuracy
loss = criterion(outputs, labels)
if phase == 'train':
loss.backward()
optimizer.step()
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
if phase == 'train':
scheduler.step() # step at end of epoch
epoch_loss = running_loss / dataset_sizes[phase]
epoch_acc = running_corrects.double() / dataset_sizes[phase]
print("{} Loss: {:.4f} Acc: {:.4f}".format(phase, epoch_loss, epoch_acc))
if phase == 'val' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = copy.deepcopy(model.state_dict()) # keep the best validation accuracy model
print()
time_elapsed = time.time() - since # slight error
print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
print("Best Val Acc: {:.4f}".format(best_acc))
model.load_state_dict(best_model_wts)
return model
After the first epoch i get this error
Epoch 0/9
----------
100%|██████████| 175/175 [15:40<00:00, 5.37s/it]
train Loss: 6.2008 Acc: 0.0037
1%|▏ | 3/232 [00:11<14:36, 3.83s/it]
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
<ipython-input-23-c552e1f19194> in <module>
----> 1 model_ft = train_model(model, criterion, optimizer, exp_lr_scheduler) # now it is a lot faster
2 # I will come back after 10 epochs
2 frames
<ipython-input-22-58764f3d79c1> in train_model(model, criterion, optimizer, scheduler, num_epochs)
26 outputs = model(inputs)
27 _, preds = torch.max(outputs, 1) # used for accuracy
---> 28 loss = criterion(outputs, labels)
29
30 if phase == 'train':
/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
1129 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130 return forward_call(*input, **kwargs)
1131 # Do not call functions when jit is used
1132 full_backward_hooks, non_full_backward_hooks = [], []
/usr/local/lib/python3.7/dist-packages/timm/loss/cross_entropy.py in forward(self, x, target)
20 def forward(self, x: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
21 logprobs = F.log_softmax(x, dim=-1)
---> 22 nll_loss = -logprobs.gather(dim=-1, index=target.unsqueeze(1))
23 nll_loss = nll_loss.squeeze(1)
24 smooth_loss = -logprobs.mean(dim=-1)
RuntimeError: CUDA error: device-side assert triggered
( my outputs are matched, my dataset has 37374 images and 488 classes
Sequential(
(0): Linear(in_features=192, out_features=150, bias=True)
(1): ReLU()
(2): Dropout(p=0.3, inplace=False)
(3): Linear(in_features=150, out_features=488, bias=True)
)
I tried os.environ['CUDA_LAUNCH_BLOCKING'] = "1" and still have the error)