My model perform this convolution operation on an image through 6 CNN pyramid levels with different resolutions. Unfortunately, the self.attenflow layer throw cuda out of memory error although the tensors have 2 channels only. could any one explain how to solve this problem 
class FlowEstimator(nn.Module):
def __init__(self, ch_in, ch_out=2, bn=False):
super(FlowEstimator, self).__init__()
self.conv1 = conv(bn, ch_in, 128)
self.conv2 = conv(bn, 128, 128)
self.conv3 = conv(bn, 128 + 128, 96)
self.conv4 = conv(bn, 128 + 96, 64)
self.conv5 = conv(bn, 96 + 64, 32)
self.attenflow=Cross_Attn(2,'relu')
self.final_out = 32
self.predict_flow = CustomConv2D(64 + 32, ch_out, kernel_size=3, stride=1,
padding=1)
def forward(self, x,context):
x1 = self.conv1(x)
x2 = self.conv2(x1)
x3 = self.conv3(torch.cat([x1, x2], dim=1))
x4 = self.conv4(torch.cat([x2, x3], dim=1))
x5 = self.conv5(torch.cat([x3, x4], dim=1))
c=context[0]
flow = self.predict_flow(torch.cat([x4, x5], dim=1))
final_flow=self.attenflow(flow,c)
return x5, final_flow