I do have all batch size equal to 1, same as that example.
Now it gives me this error:
TypeError: Broadcast function not implemented for CPU tensors
Even though I make sure all input into model are type cuda in GPU mode, it still gives me above error. It works in CPU mode, in 1 GPU mode, but not in 4 GPU mode.