code: #!/usr/bin/env python
  # -*- coding: utf-8 -*-
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import collections
  import torch
  import horovod.torch as hvd
  
  
  def broadcast_optimizer_state(optimizer, root_rank):
      """
      This function is copied from the newest horovod version.
      But the newest version has to be compiled with gcc7
      """
      if isinstance(optimizer, torch.optim.LBFGS):
          # TODO(travis): L-BFGS cannot be easily supported without serializing
          # the entire state_dict, as its structure is deeply nested and contains
          # None type parameter values
          raise ValueError('cannot broadcast torch.optim.LBFGS state')
  
      state_dict = optimizer.state_dict()
  
      # Newly created optimizers will not have their state initialized, so
      # do that initialization here
      if len(state_dict['state']) == 0:
          for group in optimizer.param_groups:
              for p in group['params']:
                  p.grad = torch.autograd.Variable(
                      p.data.new(p.size()).zero_())
          optimizer.step()
          state_dict = optimizer.state_dict()
  
      params = []
      callbacks = {}
      occurrences = collections.defaultdict(int)
  
      # Some optimizer parameters may be represented as scalars instead of
      # tensors.  In such cases, we need to wrap the scalar in a tensor, then
      # broadcast, then update the appropriate value in the state_dict with the
      # new unwrapped scalar value via a callback.
      def _create_callback(pid, name, t, p):
          def _from_tensor():
              state_dict['state'][pid][name] = t(p.numpy()[0])
          return _from_tensor
  
      # Groups are unordered, but their params will be distinct
      for group in state_dict['param_groups']:
          # The params list here is ordered by the layers in the model
          for pid in group['params']:
              if pid not in state_dict['state']:
                  continue
              param_state = state_dict['state'][pid]
              for name, p in param_state.items():
                  # Some parameter names may appear more than once, in which
                  # case we ensure they have a unique identifier defined by
                  # their order
                  occurrences[name] += 1
                  key = '%s.%d' % (str(name), occurrences[name])
  
                  if not torch.is_tensor(p):
                      # Wrap the scalar in a FloatTensor, and remember its type
                      # so we can cast it back after unwrapping
                      t = type(p)
                      p = torch.Tensor([p])
                      callbacks[key] = _create_callback(pid, name, t, p)
  
                  params.append((key, p))
  
      # Synchronized broadcast of all parameters
      hvd.broadcast_parameters(params, root_rank)
  
      # Post-broadcast clenaup for non-tensor parameters
      for key, p in params:
          if key in callbacks:
              callbacks[key]()
  
apis: torch.is_tensor torch.Tensor