code: def setup_rng(self):
      # We need to set different rng seed for different GPUs/ different runs;
      # otherwise, the noise map and time will be exactly the same.
      self.rng = torch.Generator(device=torch.cuda.current_device())
      self.rng_seed = (
          int(datetime.now().timestamp())
          + self.cfg.seed
          + parallel_state.get_data_parallel_rank()
      )
      logging.info(
          f"RNG seed set as {self.rng_seed} for rank {parallel_state.get_data_parallel_rank()}"
      )
      self.rng.manual_seed(self.rng_seed)
      self.model.set_rng(self.rng)
  
imports: import itertools
  from datetime import datetime
  from functools import partial
  from typing import Any
  import torch
  from omegaconf import DictConfig, open_dict
  from pytorch_lightning import Trainer
  from nemo.collections.multimodal.data.imagen.imagen_dataset import build_train_valid_datasets
  from nemo.collections.multimodal.models.text_to_image.imagen.precond import ContinousDDPMPrecond, EDMPrecond
  from nemo.collections.multimodal.modules.imagen.diffusionmodules.nets import EfficientUNetModel, UNetModel
  from nemo.collections.multimodal.modules.imagen.encoder.t5encoder import T5Encoder
  from nemo.collections.multimodal.modules.imagen.sampler.sampler import DDPMSampler, EDMSampler
  from nemo.collections.multimodal.parts.imagen.utils import random_dropout
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.core.classes.common import Serialization
  from nemo.utils import logging
label_desc: Use image node_id 2 for calling the Imagen obj's underlying member method code with example usage: obj.setup_rng() without return types