I cannot access the training code, but apparently the device memory is not sufficient for your current training.
Do you get this error in the first iteration or after a few iterations or epochs?
In the latter case, do you see an increase in the device memory?
If so, make sure you are not storing tensors, which are attached to the computation graph, such as the loss or output of your model.
If you are storing these tensors in e.g. a list, detach them before storing via:
losses.append(loss.detach())