The git repo is not functional. It is throwing an error. As @ptrblck mentioned the model expects an input of shape [batch_size, channels, height, width] as compared to [batch_size, height, width, channels], which is what you have given. To address this issue you would need to change the shapes of the input which can be done using torch’s permute method as shown below.
import torch
batch_size = 16
img_h,img_w = 64,64
num_channels = 3
img_1 = torch.randn(batch_size , img_h , img_w , num_channels) ### Of the form
img_1 = img_1.permute(0,3,1,2) ### Of the form [batch_size, channels, height, width]

In your case, you would need to permute x as soon as it comes into the forward method of the model.
Additionally, in your Conv2d declaration, the number of input channels is not the height. For instance, in the above example, the number of input channels is 3 and not 64. Therefore the Conv2d declaration would take an in_shape of 3.