you can leave this as with data_shape only, this does not affect the results
data_shape if batch_size!=data_shape[0] else None -> data_shape
BTW
Did your loss went to -inf?
you can leave this as with data_shape only, this does not affect the results
data_shape if batch_size!=data_shape[0] else None -> data_shape
BTW
Did your loss went to -inf?