I inserted

def forward(self, state):
    # no batch normalization if there is no batch
    print(state.shape)
    no_batch = True if len(state[0].shape) == 0 else False
    x = self.dropout_layer_1(state)


and got (while running on 4 devices):

torch.Size([72])
torch.Size([72])
torch.Size([72])
torch.Size([72])

After this it crashed