Push the model to the device via model.to('cuda') and check the memory usage via nvidia-smi.
If the model doesn’t fit, PyTorch will raise an error.
Also, since the intermediate activations will use additional memory during training, you might run a dummy forward and backward pass and check the memory usage via:
torch.cuda.memory_allocated()
torch.cuda.max_memory_allocated()