I have build exactly same model in both TF and Pytorch. And I trained in TF. For some reason, I have to transfer the pretrained weight to Pytorch.
The network is like:

image.png515×557 38.7 KB

In TF, Conv2d filter shape is [filter_height, filter_width, in_channels, out_channels], while in Pytorch is (out_channels, in_channels, kernel_size[0], kernel_size[1]).
So I have done below in TF:

and I transfer to pytorch like:

It turns out that the DQN in pytorch is not working well as in TF!