BufferError: Existing exports of data: object cannot be re-sized

[mhenn44]

I am brand new to pytorch lightning and am trying to just check that I can run pytorch lightning as intended before moving forward with anything. I have attempted two simple model codes from github and the pytorch lightning docs. Links to full code provided below and the github code has been reproduced below for ease of access

pytroch lightning docs : https://lightning.ai/docs/pytorch/stable/model/train_model_basic.html

simple model (github): https://github.com/Lightning-AI/lightning

import torch, torch.nn as nn, torch.utils.data as data, torchvision as tv, torch.nn.functional as F
import lightning as L

# Step 1: Define a LightningModule
class LitAutoEncoder(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3))
        self.decoder = nn.Sequential(nn.Linear(3, 128), nn.ReLU(), nn.Linear(128, 28 * 28))

    def forward(self, x):
        embedding = self.encoder(x)
        return embedding

    def training_step(self, batch, batch_idx):
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log("train_loss", loss)
        return loss

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

# Step 2: Define data
dataset = tv.datasets.MNIST(".", download=True, transform=tv.transforms.ToTensor())
train, val = data.random_split(dataset, [55000, 5000])


# Step 3: Train
autoencoder = LitAutoEncoder()
trainer = L.Trainer()
trainer.fit(autoencoder, data.DataLoader(train), data.DataLoader(val))

---

I've been running both in Jupyter notebook and the training just stalls with no visible error. However when I look at my terminal I see BufferError: Existing exports of data: object cannot be re-sized

Does anyone have any experience with this? As I am new I suspect this is a very beginner question but I cannot seem to find a concrete answer to this within the context of pytorch lightning specifically.

torchvision v 0.15.2a0
lightning.pytorch v 2.0.8
CPU (no GPU available)

Thank you in advance for your help

[Quantum1921]

@pytorch
@quantumsolarsevices

#QuantumDominator