original_docstring: 
      Wrapper function for `for safely loading models where certain new attributes of the model class are not present.
      E.g. "requires_stress" for computing the stress tensor.
  
      Args:
          model_path (str): Path to the model file.
          map_location (torch.device): Device where the model should be loaded to.
  
      Returns:
          :class:`schnetpack.atomistic.AtomisticModel`: Loaded SchNetPack model.
  
      
docstring: Wrapper function for `for safely loading models where certain new attributes of the model class are not present.
docstring_tokens: Wrapper function for ` for safely loading models where certain new attributes of the model class are not present .
code: def load_model(model_path, map_location=None):
      """
      Wrapper function for `for safely loading models where certain new attributes of the model class are not present.
      E.g. "requires_stress" for computing the stress tensor.
  
      Args:
          model_path (str): Path to the model file.
          map_location (torch.device): Device where the model should be loaded to.
  
      Returns:
          :class:`schnetpack.atomistic.AtomisticModel`: Loaded SchNetPack model.
  
      """
      model = torch.load(model_path, map_location=map_location)
  
      # Check for data parallel models
      if hasattr(model, "module"):
          model_module = model.module
          output_modules = model.module.output_modules
      else:
          model_module = model
          output_modules = model.output_modules
  
      # Set stress tensor attribute if not present
      if not hasattr(model_module, "requires_stress"):
          model_module.requires_stress = False
          for module in output_modules:
              module.stress = None
  
      return model
code_tokens: def load_model ( model_path , map_location = None ) : model = torch . load ( model_path , map_location = map_location ) if hasattr ( model , "module" ) : model_module = model . module output_modules = model . module . output_modules else : model_module = model output_modules = model . output_modules if not hasattr ( model_module , "requires_stress" ) : model_module . requires_stress = False for module in output_modules : module . stress = None return model
short_docstring: Wrapper function for `for safely loading models where certain new attributes of the model class are not present.
short_docstring_tokens: Wrapper function for ` for safely loading models where certain new attributes of the model class are not present .
comment: """
      Wrapper function for `for safely loading models where certain new attributes of the model class are not present.
      E.g. "requires_stress" for computing the stress tensor.
  
      Args:
          model_path (str): Path to the model file.
          map_location (torch.device): Device where the model should be loaded to.
  
      Returns:
          :class:`schnetpack.atomistic.AtomisticModel`: Loaded SchNetPack model.
  
      """$SEP$# Check for data parallel models$SEP$# Set stress tensor attribute if not present