It looks like you are directly appending the training loss to train_loss[i+1], which might hold a reference to the computation graph. If that’s the case, you are storing the computation graph in each epoch, which will grow your memory. You need to detach the loss from the computation, so that the graph can be cleared.
Change this line of code to:
train_loss[i+1] = cost_func(pred_log_probs.permute([0, 2, 1]), y_train[0, :]).detach().item()

and try it again.