This is my forward function:

def forward(self, state):
# no batch normalization if there is no batch
no_batch = True if len(state[0].shape) == 0 else False
x = self.dropout_layer_1(state)
x = F.relu(self.input_layer(x))
x = self.batch_norm_1(x) if not no_batch else x
x = self.dropout_layer_2(x)
x = F.relu(self.hidden_layer_1(x))
x = self.batch_norm_2(x) if not no_batch else x
x = self.dropout_layer_3(x)
x = F.relu(self.hidden_layer_2(x))
x = self.batch_norm_3(x) if not no_batch else x
x = self.dropout_layer_4(x)
x = self.output_layer(x)
# TODO: clamp between useful values
return x
Why could this lead to a problem? Especially if it works on CPU/ one device?

Thanks!