code: def _train_embedding(
      self,
      res,
      epochs: int,
      batch_size: int,
      lr: float,
      sample_size: int,
      num_steps: int,
      device,
  ) -> Plottable:
      _, torch, nn, _, _, _, _, trange = lazy_embed_import_dep()
      log("Training embedding")
      model, g_dataloader = res._init_model(
          res, batch_size, sample_size, num_steps, device
      )
      if (
          hasattr(res, "_embed_model")
          and not res._build_new_embedding_model
      ):
          model = res._embed_model
          log("--Reusing previous model")
  
      optimizer = torch.optim.Adam(model.parameters(), lr=lr)
      pbar = trange(epochs, desc=None)
      model.to(device)
  
      score = 0
      for epoch in pbar:
          model.train()
          for data in g_dataloader:
              g, edges, labels = data
  
              g = g.to(device)
              edges = edges.to(device)
              labels = labels.to(device)
  
              emb = model(g)
              loss = model.loss(emb, edges, labels)
              optimizer.zero_grad()
              loss.backward()
              nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
              optimizer.step()
              pbar.set_description(
                  f"epoch: {epoch+1}, loss: {loss.item():.4f}, score: {100*score:.4f}%"
              )
  
          model.eval()
          res._kg_embeddings = model(res._kg_dgl.to(device)).detach()
          res._embed_model = model
          if res._eval_flag and res._train_idx is not None:
              score = res._eval(threshold=0.5)
              pbar.set_description(
                  f"epoch: {epoch+1}, loss: {loss.item():.4f}, score: {100*score:.2f}%"
              )
  
      return res
  
imports: import logging
  import numpy
  import pandas
  from typing import Optional, Union, Callable, List, TYPE_CHECKING, Any, Tuple
  from .PlotterBase import Plottable
  from .compute.ComputeMixin import ComputeMixin
label_desc: Use image node_id 5 for calling the HeterographEmbedModuleMixin obj's underlying member method code with example usage: obj._train_embedding(res, epochs, batch_size, lr, sample_size, num_steps, device) and returns: res