RupertP:

The model works fine if the task_weights is not backpropagated and stays constant. The cuda error only appears when I set requires_grad_(True) .


Enabling the gradient calculation for additional parameters might store more intermediate activation tensors, if they are needed to compute the gradients, which could yield the out of memory issue.
You could either try to reduce the batch size or use e.g. torch.utils.checkpoint to trade compute for memory.