One subtle difference is that “_metadata” field (if any) is handled separately. See:

  

      github.com
  

  
    pytorch/pytorch/blob/e000dfcf976454fdadfdc556248976e6e560d155/torch/nn/modules/utils.py#L63 2



      
      
              state_dict (OrderedDict): a state-dict to be loaded to the model.
              prefix (str): prefix.
          """
          keys = sorted(state_dict.keys())
          for key in keys:
              if key.startswith(prefix):
                  newkey = key[len(prefix) :]
                  state_dict[newkey] = state_dict.pop(key)
          
          
# also strip the prefix in metadata if any.
          if "_metadata" in state_dict:
              metadata = state_dict["_metadata"]
              for key in list(metadata.keys()):
                  # for the metadata dict, the key can be:
                  # '': for the DDP module, which we want to remove.
                  # 'module': for the actual model.
                  # 'module.xx.xx': for the rest.
          
          
        if len(key) == 0:
                      continue
                  newkey = key[len(prefix) :]
      
    
  


  

  
    
    
  

  


Other than that, I don’t think there is a big difference. Your own implementation is a little less memory efficient I will say, as you don’t do it in place.