code: def train(
      args, model, device, train_loader, optimizer, epoch, use_fp8
  ):
      """Training function."""
      model.train()
      for batch_idx, (data, target) in enumerate(train_loader):
          data, target = data.to(device), target.to(device)
          optimizer.zero_grad()
          with te.fp8_autocast(enabled=use_fp8):
              output = model(data)
          loss = F.nll_loss(output, target)
          loss.backward()
          optimizer.step()
          if batch_idx % args.log_interval == 0:
              print(
                  f"Train Epoch: {epoch} "
                  f"[{batch_idx * len(data)}/{len(train_loader.dataset)} "
                  f"({100. * batch_idx / len(train_loader):.0f}%)]\t"
                  f"Loss: {loss.item():.6f}"
              )
              if args.dry_run:
                  break
  
file_path: ColossalAI/examples/community/fp8/mnist/main.py
imports: import argparse
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from torch.optim.lr_scheduler import StepLR
  from torchvision import datasets, transforms
label_desc: Use image node_id 1 for calling a global function with example usage: train(args, model, device, train_loader, optimizer, epoch, use_fp8) without return types