Deep Ensemble. Error in Backwards Pass

[cemde]

I am implementing a deep ensemble in lightning and keep getting an error during manual back propagation.

Let k be the number of ensemble methods. My strategy is:
1 My LightningModule has k networks
2 My DataModule has k DataLoaders
3 My Optimizer has k parameter groups
4 Each training step I get k batches and do the forward pass through the k nets.

I am not sure how to do backpropagation now. I am attempting as follows below (k=8):

class DeepEnsemble(GeneralModel):
    def __init__(self, cfg):
        super().__init__(cfg)
        self.net = nn.ModuleList([resnet18() for _ in range(8)])
        self.automatic_optimization = False
        self.temperature = nn.ModuleList([torch.nn.Parameter(torch.ones(1)) for _ in range(8)])

    def configure_optimizers(self):
        return torch.optim.Adam([{'params': net.parameters() for net in self.net}], lr=1e-5)

    def forward(self, x, idx = None):
        if idx is None:
            x = torch.stack([net.forward(x) / temperature for net, temperature in zip(self.net, self.temperature)], dim=-1)
        else:
            x = self.net[idx].forward(x) / self.temperature[idx]
        return x

    def training_step(self, batch, batch_idx):
        image, label = [b["image"] for b in batch], [b["label"] for b in batch]
        logits = [self.forward(img, i) for i, img in enumerate(image)]
        
        #loss
        loss = [self.criterion(lg, lab) for lg, lab in zip(logits, label)]
        mean_loss = torch.stack(loss).detach().mean() # detach to make sure no differentiation is done across all
        self.log('train_loss', mean_loss)

        # manual optimization
        opt = self.optimizers()
        for single_loss in loss:
            opt.zero_grad()
            self.manual_backward(single_loss)
            opt.step()
        
        # metrics
        logits_cat = torch.cat(logits, dim=0).detach()
        label_cat = torch.cat(label, dim=0)
        probs_cat = torch.softmax(logits_cat, dim=-1)
        metrics = self.log_metrics(probs_cat, label_cat, 'train')

        return None

I keep getting this error message:

Exception has occurred: RuntimeError

one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [512, 10]], which is output 0 of TBackward, is at version 10; expected version 3 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/torch/autograd/__init__.py", line 145, in backward
    Variable._execution_engine.run_backward(
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/torch/tensor.py", line 245, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 81, in backward
    closure_loss.backward(*args, **kwargs)
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 308, in backward
    output = self.precision_plugin.backward(
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 867, in backward
    self.trainer.accelerator.backward(result, optimizer, opt_idx, should_accumulate, *args, **kwargs)
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py", line 1252, in manual_backward
    self.trainer.train_loop.backward(loss, optimizer=None, opt_idx=None, *args, **kwargs)
  File "/users-1/cornelius/MislabelUnc/utils/methods.py", line 252, in training_step
    self.manual_backward(single_loss)
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/overrides/base.py", line 46, in forward
    output = self.module.training_step(*inputs, **kwargs)
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 705, in forward
    output = self.module(*inputs[0], **kwargs[0])
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp.py", line 337, in training_step
    return self.model(*args, **kwargs)
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 204, in training_step
    return self.training_type_plugin.training_step(*args)
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 290, in training_step
    training_step_output = self.trainer.accelerator.training_step(args)
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 743, in run_training_batch
    self._curr_step_result = self.training_step(
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 499, in run_training_epoch
    batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx)
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 871, in run_train
    self.train_loop.run_training_epoch()
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 809, in run_stage
    return self.run_train()
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 144, in start_training
    self._results = trainer.run_stage()
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 96, in start_training
    self.training_type_plugin.start_training(trainer)
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 799, in dispatch
    self.accelerator.start_training(self)
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 758, in _run
    self.dispatch()
  File "/users-1/cornelius/anaconda3/envs/WSS/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 460, in fit
    self._run(model)
  File "/users-1/cornelius/Project/train.py", line 53, in main
    trainer.fit(network, datamodule=datasetmodule)

What am I doing wrong? How am I supposed to do the back propagation with k ensemble members?