on_post_move_to_device is leaky
The decorator here is leaky: 
  
    
      pytorch-lightning/pytorch_lightning/core/decorators.py
    
    
        Lines 73 to 108
      in
      963c267
    
  
  
    

        
          
           def parameter_validation(fn: Callable) -> Callable: 
        

        
          
               """ 
        

        
          
               Validates that the module parameter lengths match after moving to the device. It is useful 
        

        
          
               when tying weights on TPU's. 
        

        
          
            
        

        
          
               Args: 
        

        
          
                   fn: ``model_to_device`` method 
        

        
          
            
        

        
          
               Note: 
        

        
          
                   TPU's require weights to be tied/shared after moving the module to the device. 
        

        
          
                   Failure to do this results in the initialization of new weights which are not tied. 
        

        
          
                   To overcome this issue, weights should be tied using the ``on_post_move_to_device`` model hook 
        

        
          
                   which is called after the module has been moved to the device. 
        

        
          
            
        

        
          
               See Also: 
        

        
          
                   - `XLA Documentation <https://github.com/pytorch/xla/blob/master/TROUBLESHOOTING.md#xla-tensor-quirks>`_ 
        

        
          
               """ 
        

        
          
            
        

        
          
               @wraps(fn) 
        

        
          
               def inner_fn(self, *args, **kwargs): 
        

        
          
                   pre_layer_count = len(list(self.model.parameters())) 
        

        
          
                   module = fn(self, *args, **kwargs) 
        

        
          
                   self.model.on_post_move_to_device() 
        

        
          
                   post_layer_count = len(list(self.model.parameters())) 
        

        
          
            
        

        
          
                   if not pre_layer_count == post_layer_count: 
        

        
          
                       rank_zero_warn( 
        

        
          
                           f"The model layers do not match after moving to the target device." 
        

        
          
                           " If your model employs weight sharing on TPU," 
        

        
          
                           " please tie your weights using the `on_post_move_to_device` model hook.\n" 
        

        
          
                           f"Layer count: [Before: {pre_layer_count} After: {post_layer_count}]" 
        

        
          
                       ) 
        

        
          
            
        

        
          
                   return module 
        

        
          
            
        

        
          
               return inner_fn 
        
    
  


It assumes it is called with the model_to_device method, and that self has access to model which implements on_post_move_to_device and has parameters defined.
The hook here: 
  
    
      pytorch-lightning/pytorch_lightning/core/hooks.py
    
    
        Lines 341 to 355
      in
      963c267
    
  
  
    

        
          
               def on_post_move_to_device(self) -> None: 
        

        
          
                   """ 
        

        
          
                   Called in the ``parameter_validation`` decorator after :meth:`~pytorch_lightning.core.LightningModule.to` 
        

        
          
                   is called. This is a good place to tie weights between modules after moving them to a device. Can be 
        

        
          
                   used when training models with weight sharing properties on TPU. 
        

        
          
            
        

        
          
                   Addresses the handling of shared weights on TPU: 
        

        
          
                   https://github.com/pytorch/xla/blob/master/TROUBLESHOOTING.md#xla-tensor-quirks 
        

        
          
            
        

        
          
                   Example:: 
        

        
          
            
        

        
          
                       def on_post_move_to_device(self): 
        

        
          
                           self.decoder.weight = self.encoder.weight 
        

        
          
            
        

        
          
                   """ 
        
    
  


is only called by the TPU backend. Is it intended to be called by other plugins?
Since this is under the core/decorators.py, one might assume that this is more general. Should this be an implementation detail of the TPU plugins instead? @kaushikb11
This hook is intended only for TPU as TPU don't support tying parameters on cpu.
However, I believe we could actually depreciate it once #8555 is implemented.