GPU out of memory error after few batches 
I am trying to train a complex model involving multiple convolutions. Since my own implementation was very slow (taking ~2 hours for an epoch which increased further after a few epochs), I tried changing my code to incorporate lightning module. With lightning, I'm getting a CUDA OOM after 1/3rd of the total no of batches being processed. Since the library handles .to and device calls internally, I'm not sure what can be the reason
Code sample
class Network(pl.LightningModule):
def __init__(self, in_ch=3, out_ch=1):
    super(Network, self).__init__()

    n1 = 64
    filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

    self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
    self.Up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

    self.conv0_0 = conv_block_nested(in_ch, filters[0], filters[0])
    self.conv1_0 = conv_block_nested(filters[0], filters[1], filters[1])
    self.conv2_0 = conv_block_nested(filters[1], filters[2], filters[2])
    self.conv3_0 = conv_block_nested(filters[2], filters[3], filters[3])
    self.conv4_0 = conv_block_nested(filters[3], filters[4], filters[4])

    self.conv0_1 = conv_block_nested(filters[0] + filters[1], filters[0], filters[0])
    self.conv1_1 = conv_block_nested(filters[1] + filters[2], filters[1], filters[1])
    self.conv2_1 = conv_block_nested(filters[2] + filters[3], filters[2], filters[2])
    self.conv3_1 = conv_block_nested(filters[3] + filters[4], filters[3], filters[3])

    self.conv0_2 = conv_block_nested(filters[0]*2 + filters[1], filters[0], filters[0])
    self.conv1_2 = conv_block_nested(filters[1]*2 + filters[2], filters[1], filters[1])
    self.conv2_2 = conv_block_nested(filters[2]*2 + filters[3], filters[2], filters[2])

    self.conv0_3 = conv_block_nested(filters[0]*3 + filters[1], filters[0], filters[0])
    self.conv1_3 = conv_block_nested(filters[1]*3 + filters[2], filters[1], filters[1])

    self.conv0_4 = conv_block_nested(filters[0]*4 + filters[1], filters[0], filters[0])

    self.final = nn.Conv2d(filters[0], out_ch, kernel_size=1)


def forward(self, x):
    #pdb.set_trace()  
    x0_0 = self.conv0_0(x)
    x1_0 = self.conv1_0(self.pool(x0_0))
    x0_1 = self.conv0_1(torch.cat([x0_0, self.Up(x1_0)], 1))

    x2_0 = self.conv2_0(self.pool(x1_0))
    x1_1 = self.conv1_1(torch.cat([x1_0, self.Up(x2_0)], 1))
    x0_2 = self.conv0_2(torch.cat([x0_0, x0_1, self.Up(x1_1)], 1))

    x3_0 = self.conv3_0(self.pool(x2_0))
    x2_1 = self.conv2_1(torch.cat([x2_0, self.Up(x3_0)], 1))
    x1_2 = self.conv1_2(torch.cat([x1_0, x1_1, self.Up(x2_1)], 1))
    x0_3 = self.conv0_3(torch.cat([x0_0, x0_1, x0_2, self.Up(x1_2)], 1))

    ##adding for padding issues
    x4_0 = self.conv4_0(self.pool(x3_0))
    x4_0_up = self.Up(x4_0)
    
    diffY = x3_0.size()[2] - x4_0_up.size()[2]
    diffX = x3_0.size()[3] - x4_0_up.size()[3]
    x4_0_up = F.pad(x4_0_up, (diffX // 2, diffX - diffX//2,
                    diffY // 2, diffY - diffY//2))


    x3_1 = self.conv3_1(torch.cat([x3_0, x4_0_up], 1))
    x2_2 = self.conv2_2(torch.cat([x2_0, x2_1, self.Up(x3_1)], 1))
    x1_3 = self.conv1_3(torch.cat([x1_0, x1_1, x1_2, self.Up(x2_2)], 1))
    x0_4 = self.conv0_4(torch.cat([x0_0, x0_1, x0_2, x0_3, self.Up(x1_3)], 1))
   
    output = self.final(x0_4)
    output= F.sigmoid(output).squeeze(1)
    return output

val_paths = [line.rstrip('\n') for line in open(config.val_data_path)]
train_loader = get_loader(config.train_data_path, 'Train')
val_loader = get_loader(val_paths, 'Val')
#model = torch.nn.DataParallel(Network())
model = Network()
checkpoint_callback = pl.callbacks.ModelCheckpoint( monitor="val_loss",mode="min",save_last=True, save_top_k=-1,verbose=False,)
trainer = pl.Trainer(gpus=[6,7], profiler=True, distributed_backend='ddp', num_sanity_val_steps=0,log_save_interval=1, checkpoint_callback=checkpoint_callback)
trainer.fit(model, train_dataloader=train_loader, val_dataloaders= val_loader)
Expected behavior
I'm using the same hardware so unable to understand the issue
Environment

PyTorch Version (e.g., 1.0):
OS (e.g., Linux): linux
How you installed PyTorch (conda, pip, source):conda
Build command you used (if compiling from source):
Python version: 3
CUDA/cuDNN version: 10.1
GPU models and configuration: TITAN RTX GPUs
Any other relevant information: