content: import warnings
  
  from sklearn import ensemble
  # This ignore the scheduler warning, see https://github.com/Lightning-AI/lightning/issues/5558
  warnings.filterwarnings("ignore", "Detected call of", UserWarning)
  
  import os
  import copy
  import pytorch_lightning as pl
  from CoTrain.config import ex
  from CoTrain.modules import CoTrainTransformerSS, CLIP
  from CoTrain.datamodules.video.multitask_datamodule import MTDataModule
  import datetime
  from pytorch_lightning.utilities.cloud_io import get_filesystem
  from pytorch_lightning.utilities.cloud_io import load as pl_load
  
  import torch
  import numpy as np
  from pytorch_lightning.strategies import DDPStrategy
  torch.manual_seed(0)
  
  
  class CustomDDPStrategy(DDPStrategy):
      def configure_ddp(self):
          super().configure_ddp()
          self._model._set_static_graph() # THIS IS THE MAGIC LINE
  
  
  def deterministic_index_select(x, dim, indices):
      """
      input_tensor: Tensor
      dim: dim
      indices: 1D tensor
      """
      tensor_transpose = torch.transpose(x, 0, dim)
      return tensor_transpose[indices].transpose(dim, 0)
  
  @ex.automain
  def main(_config):
      _config = copy.deepcopy(_config)
      pl.seed_everything(_config["seed"])
  
      dm = MTDataModule(_config, dist=True)
      if not _config["clip"]:
          model = CoTrainTransformerSS(_config)
      else:
          model = CLIP(_config)
      
      # assert False, [n for n, p in model.named_parameters() if p.requires_grad][:10]
  
      exp_name = f'{_config["exp_name"]}'
  
      os.makedirs(_config["log_dir"], exist_ok=True)
      checkpoint_callback = pl.callbacks.ModelCheckpoint(
          save_top_k=_config["save_top_k"],
          # every_n_epochs=_config["save_checkpoints_interval"],
          every_n_train_steps=_config["val_check_interval"],
          verbose=True,
          monitor="contrastive/train/loss",
          mode="min",
          save_last=_config["save_last"],
          dirpath=_config["model_dir"],
      )
      now = datetime.datetime.now()
      if not isinstance(_config["load_path"], str):
          instance_name = f'{exp_name}_seed{_config["seed"]}_from_multiple'
      else:
          instance_name = f'{exp_name}_seed{_config["seed"]}_from_{"_".join(_config["load_path"].split("/")[-2:])[:-5]}'
      logger = pl.loggers.TensorBoardLogger(
          _config["log_dir"],
          name=instance_name,
          version="version_0",
      )
  
      lr_callback = pl.callbacks.LearningRateMonitor(logging_interval="step")
  
      summary_callback = pl.callbacks.ModelSummary(max_depth=1)
  
      callbacks = [checkpoint_callback, lr_callback, summary_callback]
  
      num_gpus = (
          _config["num_gpus"]
          if isinstance(_config["num_gpus"], int)
          else len(_config["num_gpus"])
      )
      # print all config at the begin
      print('='*70+'Config: '+'='*70)
      print(instance_name)
      print(_config)
      print('='*150)
  
      # notice _config["batch_size"] should be max length for all machines, eg. at least 1024
      grad_steps = _config["batch_size"] // (
          _config["per_gpu_batchsize"] * num_gpus * _config["num_nodes"]
      )
      assert grad_steps > 0, (_config["batch_size"], _config["per_gpu_batchsize"])
  
      if not _config["clip_use_checkpoint"]:
          # assert not _config["clip_use_checkpoint"], "Do not use gradient accumulation and checkpoint at the same time"
          if _config["loss_names"]["openend_vqa"] >= 1:
              find_unused_paramters = True
          else:
              find_unused_paramters = False
          strategy = DDPStrategy(find_unused_parameters=find_unused_paramters)
      else:
          assert grad_steps == 1
          strategy = CustomDDPStrategy()
  
      max_steps = _config["max_steps"] if _config["max_steps"] is not None else None
      resume_ckpt = _config["resume_from"]
  
      if max_steps == None:
          max_steps = -1
  
      trainer = pl.Trainer(
          devices=_config["num_gpus"],
          num_nodes=_config["num_nodes"],
          precision=_config["precision"],
          accelerator="gpu",
          benchmark=True,
          deterministic=False,
          max_epochs=_config["max_epoch"] if max_steps == -1 else 100,
          max_steps=max_steps,
          callbacks=callbacks,
          logger=logger,
          # prepare_data_per_node=False,
          replace_sampler_ddp=False,
          accumulate_grad_batches=grad_steps,
          log_every_n_steps=10,
          fast_dev_run=_config["fast_dev_run"],
          val_check_interval=_config["val_check_interval"],
          strategy=strategy,
          # show_progress_bar=False,
          # progress_bar_refresh_rate=0
      )
  
      fs = get_filesystem(resume_ckpt)
      if fs.exists(resume_ckpt):
          with fs.open(resume_ckpt, "rb") as f:
              ckpt = torch.load(f, map_location='cpu')
          # This hacks pl wrong steps for logger
          global_step_offset = ckpt["global_step"]
          trainer.fit_loop.epoch_loop._batches_that_stepped = global_step_offset
          del ckpt
          pass
      else:
          resume_ckpt = None
  
      print("accumulate grad batches is: ", trainer.accumulate_grad_batches)
  
      if not _config["test_only"]:
          with torch.autograd.set_detect_anomaly(True):
              trainer.fit(model, datamodule=dm, ckpt_path=resume_ckpt)
      else:
          trainer.test(model, datamodule=dm, ckpt_path=resume_ckpt)
  