code: def train(
      args,
      train_g,
      sku_info,
      num_skus,
      num_brands,
      num_shops,
      num_cates,
  ):
      sampler = Sampler(
          train_g,
          args.walk_length,
          args.num_walks,
          args.window_size,
          args.num_negative,
      )
      # for each node in the graph, we sample pos and neg
      # pairs for it, and feed these sampled pairs into the model.
      # (nodes in the graph are of course batched before sampling)
      dataloader = DataLoader(
          th.arange(train_g.num_nodes()),
          # this is the batch_size of input nodes
          batch_size=args.batch_size,
          shuffle=True,
          collate_fn=lambda x: sampler.sample(x, sku_info),
      )
      model = EGES(args.dim, num_skus, num_brands, num_shops, num_cates)
      optimizer = optim.Adam(model.parameters(), lr=args.lr)
  
      for epoch in range(args.epochs):
          epoch_total_loss = 0
          for step, (srcs, dsts, labels) in enumerate(dataloader):
              # the batch size of output pairs is unfixed
              # TODO: shuffle the triples?
              srcs_embeds, dsts_embeds = model(srcs, dsts)
              loss = model.loss(srcs_embeds, dsts_embeds, labels)
  
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
  
              epoch_total_loss += loss.item()
  
              if step % args.log_every == 0:
                  print(
                      "Epoch {:05d} | Step {:05d} | Step Loss {:.4f} | Epoch Avg Loss: {:.4f}".format(
                          epoch,
                          step,
                          loss.item(),
                          epoch_total_loss / (step + 1),
                      )
                  )
  
          eval(model, test_g, sku_info)
  
      return model
  
file_path: dgl/examples/pytorch/eges/main.py
imports: import dgl
  import torch
  import torch.optim
  import utils
  from model import EGES
  from sampler import Sampler
  from sklearn import metrics
  from torch.utils.data import DataLoader
label_desc: Use image node_id 1 for calling a global function with example usage: train(args, train_g, sku_info, num_skus, num_brands, num_shops, num_cates) and returns: model