Lightning giving out of CUDA error

implementation help
1

Hi everyone,
I recently tried Pytorch lightning. I converted by pytorch code to pytorch lightning but it would give out of cuda error after few iterations where as the same code ran fine when I used pytorch. What could be the possible mistake I am making here??

2

Hi @Chetan_Pandey, would it be possible to provide a reproducible code?

3

hey @Chetan_Pandey, as @aniketmaurya suggested, i’d be great to get a reproducible script to check the issue.

We have moved the discussions to GitHub Discussions. You might want to check that out instead to get a quick response. The forums will be marked read-only after some time.

Thank you

4 
class MNISTTrainer(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_efficientnet_b4', pretrained=True)
        self.model.fc = NAC(2, 2048, 1000, 28)
        self.loss = nn.CrossEntropyLoss()
        
    def forward(self, X_batch):
        preds = self.model(X_batch)
        return preds
    
    def training_step(self, batch, batch_idx):
        img = batch['image']
        target = batch['target']
        preds = self.model(img)
        
        loss_val = self.loss(preds, target)
        self.log("Train Loss : ", loss_val)
        
        return loss_val
    
    def validation_step(self, batch, batch_idx):
        img = batch['image']
        target = batch['target']
        preds = self.model(img)
        
        loss_val = self.loss(preds, target)
        self.log("Validation Loss : ", loss_val)
        
        return loss_val
        
    def test_step(self, batch, batch_idx):
        img = batch['image'].float()
        target = batch['target']
        preds = self.model(img)
        
        loss_val = self.loss(preds, target)
        self.log("Test Loss : ", loss_val)
        
        return loss_val
    
    def predict_step(self, batch, batch_idx):
        img = batch['image']
        target = batch['target']
        
        return self.model(img)
    
    def configure_optimizers(self):
        optimizer = Adam(self.model.parameters(), lr=5e-4)
        return optimizer

classifier = MNISTTrainer()

trainer = pl.Trainer(max_epochs = 30, log_every_n_steps=20, gpus=1)

trainer.fit(classifier, train_dataloader, val_dataloader)

Sorry for late reply :sweat_smile: