Yes, I agree.
I guess the issue was caused by calling the following line every iteration
torch.nn.utils.clip_grad_norm(model.parameters(), config.clip)
Instead, I tried with another way.
First, call
params = model.parameters()
before training, and then call
torch.nn.utils.clip_grad_norm(params, config.clip)
in every iteration. Right now, the training speed is reasonable now compared to my torch implementation.