GPU usage of model decreasing each epoch
Describe the bug
Hi, i'm trying to train a model on my GPU (Tesla K80) but i have an issue : the training starts well and ~45% of the GPU is used, but at each epoch training uses less and less of the GPU slowing the training process (each epoch last longer than the previous one).
This is my Module :
class TestModule(pl.LightningModule):
    
    def __init__(self,train_dataset,dev_dataset,hparams):
        super(TestModule,self).__init__()
         
        
        self.linear1=torch.nn.Linear(2048,768)
        self.droupout=torch.nn.Dropout(p=0.1)
        self.rrelu=torch.nn.RReLU()
        self.linear2=torch.nn.Linear(768,1)
        
        self.hparams=hparams
        self.train_dataset=train_dataset
        self.dev_dataset=dev_dataset
        self.loss=torch.nn.BCEWithLogitsLoss()
        
    def forward(self,input_embedding_1,input_embedding_2):
        
       
        pooled_output_1 = input_embedding_1
        pooled_output_2 = input_embedding_2
        concatened=torch.cat((pooled_output_1,pooled_output_2),dim=1)
        logits=self.linear2(self.rrelu(self.droupout(self.linear1(concatened))))
        
        
        
        #print(act)
        
        return logits.view(-1)
    
    def training_step(self, batch, batch_nb):
        # REQUIRED
        questions_embeddings, responses_embeddings, y = batch
        y_hat = self.forward(questions_embeddings,responses_embeddings)
        loss = self.loss(y_hat, y)
        tensorboard_logs = {'train_loss': loss}
        return {'loss': loss, 'log': tensorboard_logs}

    def validation_step(self, batch, batch_nb):
        # OPTIONAL
        questions_embeddings, responses_embeddings, y = batch
        with torch.no_grad():
            y_hat = self.forward(questions_embeddings,responses_embeddings)
        predicted_labels=(y_hat>=0).float()
       
        val_acc = torch.sum(y == predicted_labels).item() / (len(y) * 1.0)
        val_acc = torch.tensor(val_acc)
        if self.on_gpu:
            val_acc = val_acc.cuda(val_acc.device.index)
        
        
        
        return {'val_loss': self.loss(y_hat, y), 'val_acc': val_acc,"predicted_labels":predicted_labels,"true_labels":y}

    def validation_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        
        avg_acc = torch.stack([x['val_acc'] for x in outputs]).mean()
        
        predictions = torch.stack([x["predicted_labels"] for x in outputs]).view(-1)
        true_labels = torch.stack([x["true_labels"] for x in outputs]).view(-1)
        
        precision = torch.tensor(((predictions[true_labels==1]==1).sum().item())/((predictions==1).sum().item()))
        recall= torch.tensor(((predictions[true_labels==1]==1).sum().item())/((true_labels==1).sum().item()))
        
        f1_score=2*((precision*recall)/(precision+recall))
                                                             
        if self.on_gpu:
            f1_score = f1_score.cuda(f1_score.device.index)
        tensorboard_logs = {'val_loss': avg_loss,"val_acc":avg_acc,"f1_score":f1_score}
        return {'progress_bar': tensorboard_logs, 'log': tensorboard_logs}

    def configure_optimizers(self):
        # REQUIRED
        # can return multiple optimizers and learning_rate schedulers
        # (LBFGS it is automatically supported, no need for closure function)
        return torch.optim.Adam(self.parameters(), lr=self.hparams.lr)

    @pl.data_loader
    def train_dataloader(self):
        # REQUIRED
        data_train_sampler = RandomSampler(self.train_dataset)        
        return DataLoader(self.train_dataset, sampler=data_train_sampler, batch_size=self.hparams.batch_size,drop_last=True,pin_memory=True)

    @pl.data_loader
    def val_dataloader(self):
        data_dev_sampler = SequentialSampler(self.dev_dataset)        
        return DataLoader(self.dev_dataset, sampler=data_dev_sampler, batch_size=self.hparams.batch_size,drop_last=True,pin_memory=True)

trainer = Trainer(gpus=1,default_save_path="output_non_linear_without_duplicated/",check_val_every_n_epoch=1,accumulate_grad_batches=2)

model = TestModule(train_dataset=train_ds,dev_dataset=dev_ds,hparams=config)

trainer.fit(model)

I've tried the same model but training with a standard PyTorch training loop and all goes well (the training uses ~60% of the GPU and the usage remains stable).
Do I have made an error building my PyTorch Lightning module ?
Expected behavior
Model uses the same amount of GPU than the same module in standard PyTorch and the GPU usage remains stable across each training epoch.
Desktop (please complete the following information):

OS: Ubuntu 18.04
Python 3.7, PyTorch 1.2