# How to train mnist with trainer?

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-train-mnist-with-trainer%2F64960%3E%5C
**Category:** 🤗Transformers\
**Created:** [December 7, 2023, 5:28pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-train-mnist-with-trainer%2F64960 "2023-12-07T17:28:53Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![Ligeng-Zhu](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/ligeng-zhu/32/21965_2.png) [@Ligeng-Zhu](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FLigeng-Zhu)%5C
**Post date:** [December 7, 2023, 5:28pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-train-mnist-with-trainer%2F64960%2F1 "2023-12-07T17:28:54Z")

</div>

Hi there,

I am trying to switch some vision training codes to trainer, however, I am bit confused about how to make the model “trainer-compatible”. For example, I take the reference from [From PyTorch DDP to Accelerate to Trainer, mastery of distributed training with ease](https://huggingface.co/blog/pytorch-ddp-accelerate-transformers)

```python
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from transformers import Trainer, TrainingArguments

class BasicNet(nn.Module):
    def __init__ (self):
        super(). __init__ ()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)
        self.act = F.relu

    def forward(self, x):
        x = self.act(self.conv1(x))
        x = self.act(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.act(self.fc1(x))
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

device = "cuda"
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307), (0.3081))
])

train_dset = datasets.MNIST('data', train=True, download=True, transform=transform)
test_dset = datasets.MNIST('data', train=False, transform=transform)

train_loader = torch.utils.data.DataLoader(train_dset, shuffle=True, batch_size=64)
test_loader = torch.utils.data.DataLoader(test_dset, shuffle=False, batch_size=64)

model = BasicNet()
training_args = TrainingArguments(
    "basic-trainer",
    per_device_train_batch_size=64,
    per_device_eval_batch_size=64,
    num_train_epochs=3,
    evaluation_strategy="epoch",
    remove_unused_columns=False
)

def collate_fn(examples):
    pixel_values = torch.stack([example[0] for example in examples])
    labels = torch.tensor([example[1] for example in examples])
    return {"x":pixel_values, "labels":labels}

class MyTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        outputs = model(inputs["x"])
        target = inputs["labels"]
        loss = F.nll_loss(outputs, target)
        return (loss, outputs) if return_outputs else loss

trainer = MyTrainer(
    model,
    training_args,
    train_dataset=train_dset,
    eval_dataset=test_dset,
    data_collator=collate_fn,
)

trainer.train()

```

This snippet will function properlly for training, however, when launcing evaluation, it shows error

```python
Original Traceback (most recent call last):
  File "/home/ligengz/anaconda3/envs/awq/lib/python3.9/site-packages/torch/nn/parallel/parallel_apply.py", line 64, in _worker
    output = module(*input, **kwargs)
  File "/home/ligengz/anaconda3/envs/awq/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
TypeError: forward() got an unexpected keyword argument 'labels'

```

I guess there are wraps we need to apply on model. What is the suggested way to convert pytorch model to hf-trainer compatible ones?

---

<div class="post-metadata">

**Author:** ![Ligeng-Zhu](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/ligeng-zhu/32/21965_2.png) [@Ligeng-Zhu](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FLigeng-Zhu)%5C
**Post date:** [December 7, 2023, 5:46pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-train-mnist-with-trainer%2F64960%2F2 "2023-12-07T17:46:45Z")

</div>

Solved.

Simply changing forward args from  
` def forward(self, x):`  
to  
` def forward(self, pixel_values, labels=None):`

then the code will work as expected.