class_docstr: Base class for distributed computation models and defines interface methods.
  This class is public and should be used for other custom derived distributed models.
class_docstr_tok: ["Base","class","for","distributed","computation","models","and","defines","interface","methods",".","This","class","is","public","and","should","be","used","for","other","custom","derived","distributed","models","."]
code: def _apply_op(
      self,
      tensor: torch.Tensor,
      device: torch.device,
      fn: Callable,
      *args: Any,
      **kwargs: Any,
  ) -> torch.Tensor:
      out_dtype = None
      tensor_device = None
  
      # check if the tensor is at specified device
      if tensor.device != device:
          tensor_device = tensor.device
          tensor = tensor.to(device)
  
      if self._collective_op_dtype is not None:
          # cast to _collective_op_dtype if current type is not floatX
          if tensor.dtype not in (torch.float32, torch.float64):
              out_dtype = tensor.dtype
              tensor = tensor.to(self._collective_op_dtype)
  
      tensor = fn(tensor, *args, **kwargs)
  
      if out_dtype is not None and tensor_device is not None:
          return tensor.to(dtype=out_dtype, device=tensor_device)
      if out_dtype is not None:
          return tensor.to(dtype=out_dtype)
      if tensor_device is not None:
          return tensor.to(device=tensor_device)
      return tensor
  
imports: from abc import ABCMeta, abstractmethod
  from numbers import Number
  from typing import Any, Callable, cast, List, Optional, Union
  import torch
  from packaging.version import Version
label_desc: Use image node_id 22 for calling the ComputationModel obj's underlying member method code with example usage: obj._apply_op(tensor, device, fn) and returns: tensor, tensor, tensor, tensor