I realise where I was making a mistake. My model has an LSTM and I’m supposed to pass on a new, empty variable as the hidden state. If I pass on an existing variable, such as the hidden state from the previous timestep, the model backprops all the way back to the first epoch on every epoch of training. This is precisely why the GPU memory kept exploding after every epoch. I now have something like this, and it works fine.
def zero_hidden():
    return (torch.zeros(1, 1, hidden_dim),
            torch.zeros(1, 1, hidden_dim))

lstm_out, lstm_hidden = lstm(lstm_in, zero_hidden())