text: import copy
  import time
  import datetime
  
  import torch
  import numpy as np
  
  from torchreid import metrics
  from torchreid.losses import TripletLoss, CrossEntropyLoss
  
  from torchreid.engine import Engine
  
  from torch.utils.tensorboard import SummaryWriter
  
  
  class ImageMutualEngine(Engine):
      def __init__(
          self,
          datamanager,
          model_1,
          model_2,
          model_cluster,
          optimizer_1,
          optimizer_2,
          margin,
          weight_t,
          weight_x,
          scheduler_1,
          scheduler_2,
          use_gpu,
          label_smooth,
      ):
          super().__init__(
              datamanager,
              use_gpu=use_gpu,
          )
  
          self.model_1 = model_1
          self.model_2 = model_2
          self.optimizer_1 = optimizer_1
          self.optimizer_2 = optimizer_2
          self.scheduler_1 = scheduler_1
          self.scheduler_2 = scheduler_2
          self.register_model("model_1", model_1, optimizer_1, scheduler_1)
          self.register_model("model_2", model_2, optimizer_2, scheduler_2)
  
          self.mean_model_1 = copy.deepcopy(self.model_1)
          self.mean_model_2 = copy.deepcopy(self.model_2)
  
          for param in self.mean_model_1.parameters():
              param.detach()
          for param in self.mean_model_2.parameters():
              param.detach()
  
          self.model_cluster = model_cluster
  
          assert weight_t >= 0 and weight_x >= 0
          assert weight_t + weight_x > 0
          self.weight_t = weight_t
          self.weight_x = weight_x
  
          # Define losses
          self.criterion_t = None
          self.criterion_st = None
          self.criterion_x = CrossEntropyLoss(
              num_classes=self.datamanager.num_train_pids,
              use_gpu=self.use_gpu,
              label_smooth=label_smooth,
          )
          self.criterion_sx = None
  
      def run(
          self,
          save_dir="log",
          max_epoch=0,
          start_epoch=0,
          print_freq=10,
          fixbase_epoch=0,
          open_layers=None,
          start_eval=0,
          eval_freq=-1,
          test_only=False,
          dist_metric="euclidean",
          normalize_feature=False,
          visrank=False,
          visrank_topk=10,
          use_metric_cuhk03=False,
          ranks=[1, 5, 10, 20],
          rerank=False,
      ):
          if visrank and not test_only:
              raise ValueError("visrank can be set to True only if test_only=True")
  
          if test_only:
              self.test(
                  dist_metric=dist_metric,
                  normalize_feature=normalize_feature,
                  visrank=visrank,
                  visrank_topk=visrank_topk,
                  save_dir=save_dir,
                  use_metric_cuhk03=use_metric_cuhk03,
                  ranks=ranks,
                  rerank=rerank,
              )
              return
  
          if self.writer is None:
              self.writer = SummaryWriter(log_dir=save_dir)
  
          time_start = time.time()
          self.start_epoch = start_epoch
          self.max_epoch = max_epoch
          print("=> Start training")
  
          for self.epoch in range(self.start_epoch, self.max_epoch):
              # Generate pseudo labels
              self._generate_pseudo_labels(self.train_loader)
              self.train(print_freq=print_freq, fixbase_epoch=fixbase_epoch, open_layers=open_layers)
  
              if (
                  (self.epoch + 1) >= start_eval
                  and eval_freq > 0
                  and (self.epoch + 1) % eval_freq == 0
                  and (self.epoch + 1) != self.max_epoch
              ):
                  rank1 = self.test(
                      dist_metric=dist_metric,
                      normalize_feature=normalize_feature,
                      visrank=visrank,
                      visrank_topk=visrank_topk,
                      save_dir=save_dir,
                      use_metric_cuhk03=use_metric_cuhk03,
                      ranks=ranks,
                  )
                  self.save_model(self.epoch, rank1, save_dir)
  
          if self.max_epoch > 0:
              print("=> Final test")
              rank1 = self.test(
                  dist_metric=dist_metric,
                  normalize_feature=normalize_feature,
                  visrank=visrank,
                  visrank_topk=visrank_topk,
                  save_dir=save_dir,
                  use_metric_cuhk03=use_metric_cuhk03,
                  ranks=ranks,
              )
              self.save_model(self.epoch, rank1, save_dir)
  
          elapsed = round(time.time() - time_start)
          elapsed = str(datetime.timedelta(seconds=elapsed))
          print("Elapsed {}".format(elapsed))
          if self.writer is not None:
              self.writer.close()
  
      def forward_backward(self, data):
          pass
  
      def _generate_pseudo_labels(self, dataloader):
          print(f"Encoding features for clustering.")
          dataloader.dataset.mutual = False
          full_features = []
          for samples, _ in dataloader:
              samples = samples.to(device)
              self.mean_model_1(samples)
              self.mean_model_2(samples)
              batch_features_1 = self.mean_model_1.hooks.numpy()
              batch_features_2 = self.mean_model_2.hooks.numpy()
              for f in zip(batch_features_1, batch_features_2):
                  full_features.append((f[0] + f[1]) / 2)
  
          full_features = np.array(full_features, dtype=object)
          pseudo_label = self.model_cluster.generate_pseudo_labels(full_features)
          dataloader.dataset.pseudo_labels = torch.tensor(pseudo_label, dtype=torch.long)
  