code: def test(self, comp_type, comm_type):
      if comp_type == torch.bfloat16 or comm_type == torch.bfloat16:
          if not bf16_required_version_check():
              pytest.skip(
                  " DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly"
              )
  
      type_str = {torch.float16: "fp16", torch.bfloat16: "bfp16"}
  
      config_dict = {
          "train_batch_size": 2,
          "steps_per_print": 1,
          "fp16": {"enabled": comp_type == torch.float16},
          "bf16": {"enabled": comp_type == torch.bfloat16},
          "zero_optimization": {"stage": 2},
          "communication_data_type": type_str[comm_type],
      }
      hidden_dim = 10
  
      model = SimpleModel(hidden_dim)
      optimizer = torch.optim.Adam(model.parameters())
      model, _, _, _ = deepspeed.initialize(
          config=config_dict, model=model, optimizer=optimizer
      )
      data_loader = random_dataloader(
          model=model,
          total_samples=2,
          hidden_dim=hidden_dim,
          device=model.device,
          dtype=comp_type,
      )
  
      def custom_reduce(
          tensor, dst, op=dist.ReduceOp.SUM, group=None, async_op=False
      ):
          assert tensor.dtype == comm_type
          return orig_torch_reduce(tensor, dst, op, group, async_op)
  
      orig_torch_reduce = dist.reduce
      dist.reduce = custom_reduce
      for n, batch in enumerate(data_loader):
          loss = model(batch[0], batch[1])
          model.backward(loss)
          model.step()
      dist.reduce = orig_torch_reduce
  
imports: import torch
  import deepspeed
  import pytest
  from deepspeed.ops.adam import FusedAdam
  from unit.common import DistributedTest
  from deepspeed.ops.op_builder import CPUAdamBuilder
  from unit.simple_model import SimpleModel, SimpleOptimizer, random_dataloader
  from unit.util import bf16_required_version_check
  from deepspeed import comm
label_desc: Use image node_id 1 for calling the TestZeroDtypeCocktail obj's underlying member method code with example usage: obj.test(comp_type, comm_type) and returns: orig_torch_reduce