Also, I think another error is in calling model.cuda(device=0).
This will move all model parameters and buffers to the GPU specified with the id.
So your self.conv2 layer will also be moved to GPU0, which might cause the same error again.
If you really need model sharding, I think you should define the splits in the __init__ function of your model.