I was able to generate an example that reproduces the problem and after some digging …
It is a reflection of the problems discussed here:
Issues In the pytorch github:

59923
35643

and it seems like a fix for kwargs to be handled correctly by the hooks is in the making

  

      github.com/pytorch/pytorch
  

  
    
  
    
  

  
    
      [pytorch] Lazy Module Fix for kwargs argument support on forward 5
    

    
      pytorch:master ← ertrue:export-D29687148
    

    
      
        opened 
        
          
        
        Jul 13, 2021
      
      

      
        
          
          ertrue
        
      

      
        
          +59
          -15
        
      
    
  


  
    Summary:
This is a fix to resolve the issue: #35643 and #59923

Test Plan:
p…ython3 test: caffe2/test/test_nn.py

Differential Revision: D29687148
  

  

  
    
    
  

  


TLDR: Keywords argument dont show up in hooks, fix your code until they do
import torch
from torch import nn


class BorinModel(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear1 = nn.Linear(in_features=10, out_features=1)
        self.linear2 = nn.Linear(in_features=20, out_features=1)

    def forward(self, x, y, debug=True):
        if debug:
            print("im debuggin")

        return self.linear1(x) + self.linear2(y)

class EvenMoreBoringModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.boring_stuff = BorinModel()

    def forward(self, x, y):
        return self.boring_stuff(x, y, debug=True)

    def forward2(self, x, y):
        return self.boring_stuff(x=x, y=y, debug=True)

def hook(m, i, o):
    print(m)
    print(f"Inputs: {[x.shape for x in i]}")
    print(f"Outputs: {o.shape}")

boring_instance = EvenMoreBoringModel()

handle = boring_instance.boring_stuff.register_forward_hook(hook)

boring_instance.forward(torch.ones(5,10), torch.ones(5, 20))
# im debuggin
# BorinModel(
#   (linear1): Linear(in_features=10, out_features=1, bias=True)
#   (linear2): Linear(in_features=20, out_features=1, bias=True)
# )
# Inputs: [torch.Size([5, 10]), torch.Size([5, 20])]
# Outputs: torch.Size([5, 1])
boring_instance.forward2(torch.ones(5,10), y = torch.ones(5, 20))
# im debuggin
# BorinModel(
#   (linear1): Linear(in_features=10, out_features=1, bias=True)
#   (linear2): Linear(in_features=20, out_features=1, bias=True)
# )
# Inputs: []
# Outputs: torch.Size([5, 1])