code: def train(args):
      logger.warning(
          "This script is an example to showcase the MetaModule and "
          "data-loading features of Torchmeta, and as such has been "
          "very lightly tested. For a better tested implementation of "
          "Model-Agnostic Meta-Learning (MAML) using Torchmeta with "
          "more features (including multi-step adaptation and "
          "different datasets), please check `https://github.com/"
          "tristandeleu/pytorch-maml`."
      )
  
      dataset = omniglot(
          args.folder,
          shots=args.num_shots,
          ways=args.num_ways,
          shuffle=True,
          test_shots=15,
          meta_train=True,
          download=args.download,
      )
      dataloader = BatchMetaDataLoader(
          dataset,
          batch_size=args.batch_size,
          shuffle=True,
          num_workers=args.num_workers,
      )
  
      model = ConvolutionalNeuralNetwork(
          1, args.num_ways, hidden_size=args.hidden_size
      )
      model.to(device=args.device)
      model.train()
      meta_optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
  
      # Training loop
      with tqdm(dataloader, total=args.num_batches) as pbar:
          for batch_idx, batch in enumerate(pbar):
              model.zero_grad()
  
              train_inputs, train_targets = batch["train"]
              train_inputs = train_inputs.to(device=args.device)
              train_targets = train_targets.to(device=args.device)
  
              test_inputs, test_targets = batch["test"]
              test_inputs = test_inputs.to(device=args.device)
              test_targets = test_targets.to(device=args.device)
  
              outer_loss = torch.tensor(0.0, device=args.device)
              accuracy = torch.tensor(0.0, device=args.device)
              for task_idx, (
                  train_input,
                  train_target,
                  test_input,
                  test_target,
              ) in enumerate(
                  zip(
                      train_inputs,
                      train_targets,
                      test_inputs,
                      test_targets,
                  )
              ):
                  train_logit = model(train_input)
                  inner_loss = F.cross_entropy(
                      train_logit, train_target
                  )
  
                  model.zero_grad()
                  params = gradient_update_parameters(
                      model,
                      inner_loss,
                      step_size=args.step_size,
                      first_order=args.first_order,
                  )
  
                  test_logit = model(test_input, params=params)
                  outer_loss += F.cross_entropy(test_logit, test_target)
  
                  with torch.no_grad():
                      accuracy += get_accuracy(test_logit, test_target)
  
              outer_loss.div_(args.batch_size)
              accuracy.div_(args.batch_size)
  
              outer_loss.backward()
              meta_optimizer.step()
  
              pbar.set_postfix(
                  accuracy="{0:.4f}".format(accuracy.item())
              )
              if batch_idx >= args.num_batches:
                  break
  
      # Save model
      if args.output_folder is not None:
          filename = os.path.join(
              args.output_folder,
              "maml_omniglot_"
              "{0}shot_{1}way.th".format(args.num_shots, args.num_ways),
          )
          with open(filename, "wb") as f:
              state_dict = model.state_dict()
              torch.save(state_dict, f)
  
file_path: pytorch-meta/examples/maml/train.py
imports: import os
  import torch
  import torch.nn.functional
  from tqdm import tqdm
  import logging
  from torchmeta.datasets.helpers import omniglot
  from torchmeta.utils.data import BatchMetaDataLoader
  from torchmeta.utils.gradient_based import gradient_update_parameters
  from model import ConvolutionalNeuralNetwork
  from utils import get_accuracy
label_desc: Use image node_id 1 for calling a global function with example usage: train(args) without return types