code: def train(dataloader, args):
      writer = SummaryWriter("runs")
      net = Net()
      net.train()
      optimizer = SGD(
          net.parameters(),
          lr=args.lr,
          momentum=args.momentum,
          weight_decay=args.wd,
      )
      gm = GradManager().attach(net.parameters())
  
      epoch_length = len(dataloader)
      for epoch in range(args.epoch):
          for step, (batch_data, batch_label) in enumerate(dataloader):
              batch_label = batch_label.astype(np.int32)
              data, label = mge.tensor(batch_data), mge.tensor(
                  batch_label
              )
              with gm:
                  pred = net(data)
                  loss = F.loss.cross_entropy(pred, label)
                  gm.backward(loss)
              optimizer.step().clear_grad()
  
              if step % 50 == 0:
                  print(
                      "epoch:{}, iter:{}, loss:{}".format(
                          epoch + 1, step, float(loss)
                      )
                  )  # noqa
              writer.add_scalar(
                  "loss", float(loss), epoch * epoch_length + step
              )
          if (epoch + 1) % 5 == 0:
              mge.save(
                  net.state_dict(),
                  os.path.join(
                      gettempdir(), f"mnist_net_e{epoch + 1}.pkl"
                  ),
              )
  
file_path: clearml/examples/frameworks/megengine/megengine_mnist.py
imports: import argparse
  import os
  import sys
  from tempfile import gettempdir
  import numpy
  from clearml import Task
  from tensorboardX import SummaryWriter
label_desc: Use image node_id 2 for calling a global function with example usage: train(dataloader, args) without return types