'UNet2DOutput' object has no attribute 'size'

Question

net = UNet2DModel(sample_size=28,
                    in_channels=1,
                    out_channels=1,
                    layers_per_block=2,
                    block_out_channels=(32, 64, 64),
                    down_block_types=(
                        "DownBlock2D",
                        "AttnDownBlock2D",
                        "AttnDownBlock2D",
                    ),
                    up_block_types=(
                        "AttnUpBlock2D",
                        "AttnUpBlock2D",
                        "UpBlock2D",
                    ),
)

batch_size = 64
n_epochs = 3
train_dataloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True)

net.to(device)

loss_fn = nn.MSELoss()

opt = torch.optim.Adam(net.parameters(), lr=1e-3)

losses = [] 

for epoch in range(n_epochs):
    for x, y in train_dataloader:
        x = x.to(device)
        noise_amount = torch.rand(x.shape[0]).to(device)
        
        noisy_x = corrupt(x, noise_amount)
        
        pred = net(noisy_x, 0)
        
        loss = loss_fn(pred, x)
        
        opt.zero_grad()
        loss.backward()
        opt.step()
        
        losses.append(loss.item())

the output is AttributeError: 'UNet2DOutput' object has no attribute 'size'. the error is occurred in " loss = loss_fn(pred, x) ". I am a begginer in torch, what can I do to solve this problem?