Typically a gradual OOM after many epochs can be the result of something in training loop unwittingly holding on to previous data that no longer needs to stored (e.g., the loss). Can you share the training loop of the model?