question: on_post_move_to_device is leaky
context: 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
answer: 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.