When ensembling a kfold trained set, does it lose the ability to set other models?

[VDFaller]

Almost definitely me doing something wrong, but I'm trying to ensemble 3 of the same backbone model. When I do the simple torch stack, mean. For whatever reason, the models it's for()ing through seem to lose the ability to set the device and eval mode. Here's the simplest example I could come up with. Am I doing something obviously wrong? I can do it if each model is an individual variable, I assume that's by design.

Edit and as soon as I press the post button. I think I found my answer. I can just add_module??

# %%
import torch
from pytorch_lightning import LightningModule, Trainer
from torchvision.datasets import MNIST
from torchvision.models import resnet18
from torchvision import transforms


# %%
class LitModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.backbone = resnet18(num_classes=10)
        self.backbone.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.backbone.fc = torch.nn.Linear(self.backbone.fc.in_features, 10)
        self.criterion = torch.nn.CrossEntropyLoss()

    def forward(self, x):
        return self.backbone(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.criterion(y_hat, y)
        self.log('train_loss', loss)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

class KFoldLitModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.models = [LitModel() for _ in range(3)] 
        self.criterion = torch.nn.CrossEntropyLoss()

    def forward_what_want_to_work(self, x):
        """This is ONLY used for the testing/inference phase becaues I can't figure out how to make good batches for kfold."""
        # when I do this, the models are not in eval mode and it's not set to the proper device
        logits = torch.stack([model(x) for model in self.models], dim=0).mean(dim=0)
        return logits

    def forward(self, x):
        device = x.device
        logits = []
        for model in self.models:
            model.eval()  # Obviously I wouldn't want to do this if I were training here.
            logits.append(model.to(device)(x)) # not doing this breaks stuff. 
        logits = torch.stack(logits, dim=0)
        logits = logits.mean(dim=0)
        return logits

    def test_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.criterion(y_hat, y)
        return loss


dataset = MNIST(root='.', download=True, transform=transforms.ToTensor())
dataloader = torch.utils.data.DataLoader(dataset, batch_size=256)

model = KFoldLitModel()
trainer = Trainer(max_epochs=1)
# doing this 
trainer.test(model, dataloaders=dataloader)