code: def main(args):
      device = torch.device(
          "cuda:0" if torch.cuda.is_available() else "cpu"
      )
  
      if args.use_deterministic_algorithms:
          torch.backends.cudnn.benchmark = False
          torch.use_deterministic_algorithms(True)
      else:
          torch.backends.cudnn.benchmark = True
  
      p = args.labels_per_batch
      k = args.samples_per_label
      batch_size = p * k
  
      model = EmbeddingNet()
      if args.resume:
          model.load_state_dict(
              torch.load(args.resume, weights_only=True)
          )
  
      model.to(device)
  
      criterion = TripletMarginLoss(margin=args.margin)
      optimizer = Adam(model.parameters(), lr=args.lr)
  
      transform = transforms.Compose(
          [
              transforms.Lambda(lambda image: image.convert("RGB")),
              transforms.Resize((224, 224)),
              transforms.PILToTensor(),
              transforms.ConvertImageDtype(torch.float),
          ]
      )
  
      # Using FMNIST to demonstrate embedding learning using triplet loss. This dataset can
      # be replaced with any classification dataset.
      train_dataset = FashionMNIST(
          args.dataset_dir,
          train=True,
          transform=transform,
          download=True,
      )
      test_dataset = FashionMNIST(
          args.dataset_dir,
          train=False,
          transform=transform,
          download=True,
      )
  
      # targets is a list where the i_th element corresponds to the label of i_th dataset element.
      # This is required for PKSampler to randomly sample from exactly p classes. You will need to
      # construct targets while building your dataset. Some datasets (such as ImageFolder) have a
      # targets attribute with the same format.
      targets = train_dataset.targets.tolist()
  
      train_loader = DataLoader(
          train_dataset,
          batch_size=batch_size,
          sampler=PKSampler(targets, p, k),
          num_workers=args.workers,
      )
      test_loader = DataLoader(
          test_dataset,
          batch_size=args.eval_batch_size,
          shuffle=False,
          num_workers=args.workers,
      )
  
      if args.test_only:
          # We disable the cudnn benchmarking because it can noticeably affect the accuracy
          torch.backends.cudnn.benchmark = False
          torch.backends.cudnn.deterministic = True
          evaluate(model, test_loader, device)
          return
  
      for epoch in range(1, args.epochs + 1):
          print("Training...")
          train_epoch(
              model,
              optimizer,
              criterion,
              train_loader,
              device,
              epoch,
              args.print_freq,
          )
  
          print("Evaluating...")
          evaluate(model, test_loader, device)
  
          print("Saving...")
          save(model, epoch, args.save_dir, "ckpt.pth")
  
file_path: vision/references/similarity/train.py
imports: import os
  import torch
  import torchvision.transforms
  from loss import TripletMarginLoss
  from model import EmbeddingNet
  from sampler import PKSampler
  from torch.optim import Adam
  from torch.utils.data import DataLoader
  from torchvision.datasets import FashionMNIST
label_desc: Use image node_id 5 for calling a global function with example usage: main(args) without return types