Olorin:

while Keras occupied all GPU memory during the training


That’s the default behavior of TensorFlow and doesn’t mean that the complete GPU memory will be used.
By skimming through your code I cannot find any obvious issues, so you could load the Keras parameters into the PyTorch model and compare the output for a static input to make sure the architecture is equal.