code: def train(args):
      logger.warning(
          "This script is an example to showcase the extensions and "
          "data-loading features of Torchmeta, and as such has been "
          "very lightly tested."
      )
  
      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 = MatchingNetwork(
          1, args.embedding_size, hidden_size=args.hidden_size
      )
      model.to(device=args.device)
      model.train()
      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)
              train_embeddings = model(train_inputs)
  
              test_inputs, test_targets = batch["test"]
              test_inputs = test_inputs.to(device=args.device)
              test_targets = test_targets.to(device=args.device)
              test_embeddings = model(test_inputs)
  
              loss = matching_loss(
                  train_embeddings,
                  train_targets,
                  test_embeddings,
                  test_targets,
                  args.num_ways,
              )
              loss.backward()
              optimizer.step()
  
              with torch.no_grad():
                  # calculate the accuracy
                  log_probas = matching_log_probas(
                      train_embeddings,
                      train_targets,
                      test_embeddings,
                      args.num_ways,
                  )
                  test_predictions = torch.argmax(log_probas, dim=1)
                  accuracy = torch.mean(
                      (test_predictions == test_targets).float()
                  )
                  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,
              "matching_network_omniglot_"
              "{0}shot_{1}way.pt".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/matching-network/train.py
imports: import os
  import torch
  from tqdm import tqdm
  import logging
  from torchmeta.datasets.helpers import omniglot
  from torchmeta.utils.data import BatchMetaDataLoader
  from torchmeta.utils.matching import matching_log_probas, matching_loss
  from model import MatchingNetwork
label_desc: Use image node_id 1 for calling a global function with example usage: train(args) without return types