I have an architecture, built to process large 3D images. During a forward pass (for batch_size=1), GPU memory overshoots at place of the 3D transposed convolution.

self.add_module(‘U_conv_’+str(i+1), nn.ConvTranspose3d(channel[i], channel[i+1], kernel_size, stride=2, padding=padd, output_padding=out_padd, bias=True)

…

File “/home/ivan/main.py”, line 149, in forward
x = self._modules’U_conv_’+str(i+1)
File “/home/ivan/.local/lib/python2.7/site-packages/torch/nn/modules/module.py”, line 477, in call
result = self.forward(*input, **kwargs)
File “/home/ivan/.local/lib/python2.7/site-packages/torch/nn/modules/conv.py”, line 818, in forward
output_padding, self.groups, self.dilation)
RuntimeError: CUDA error: out of memory

Is there a way to parallelize over multiple GPUs only a single layer (due to large size of the image, parallelizing different layers didn’t solve the issue)?
Thanks!