Pythonでモデルのトレーニング中に「AttributeError: 'tuple' object has no attribute 'log

Pythonでモデルをトレーニングしようとしているのですが、AttributeError: 'tuple' object has no attribute 'log_softmax'というエラーが発生してしまいます。
対処方法を教えていただきたいです。

発生している問題・エラーメッセージ

エラーメッセージ
Traceback (most recent call last):
  File "models_mae.py", line 367, in <module>
    train(epoch)
  File "models_mae.py", line 343, in train
    loss = criterion(output, target) 
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/loss.py", line 961, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/functional.py", line 2468, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/functional.py", line 1605, in log_softmax
    ret = input.log_softmax(dim)
AttributeError: 'tuple' object has no attribute 'log_softmax'


### 該当のソースコード

（モデルの定義…字数の関係で省略）

def mae_vit_huge_patch14_dec512d8b(**kwargs):
    model = MaskedAutoencoderViT(
        patch_size=14, embed_dim=1280, depth=32, num_heads=16,
        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model

mae_vit_large_patch16 = mae_vit_large_patch16_dec512d8b 


model = mae_vit_large_patch16()
epoch_count = 10

device = torch.device('cuda')
model = model.to(device)

import torchvision.transforms as transforms
import torchvision.models as models
import torchvision

torch.manual_seed(0)
transform = transforms.Compose(
            [
            # 画像サイズが異る場合は利用して画像サイズを揃える
            transforms.Resize((224,224)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor()
             ])

data = torchvision.datasets.ImageFolder(root='/home/img', transform=transform)

train_size = int(0.8 * len(data))
validation_size  = int(0.1 * len(data))
test_size= int(0.1 * len(data))
data_size  = {"train":train_size, "validation":validation_size,"test_size":test_size}

data_train, data_validation,data_test = torch.utils.data.random_split(data, [train_size, validation_size,test_size], generator=torch.Generator().manual_seed(42))

train_loader = torch.utils.data.DataLoader(data_train, batch_size=16, shuffle=True)

import torch.nn as nn
import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

def train(epoch):
    total_loss = 0 
    total_size = 0 
    model.train() 
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        total_loss += loss.item()
        total_size += data.size(0)
        loss.backward()
        optimizer.step()
        
        average_loss = total_loss / total_size
        
        loss_list=[] 
        loss_list.append(average_loss) 
        
        if batch_idx % 1000 == 0: 
            now = datetime.datetime.now()
            print('[{}] Train Epoch: {} [{}/{} ({:.0f}%)]\tAverage loss: {:.6f}'.format(
                now,
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), total_loss / total_size))