Solved by ptrblck in post #2 
              
              
                module1 should create valid gradients as long as mask2 isn’t full of zeros as seen here: 
module1 = nn.Linear(10, 10)
module2 = nn.Linear(10, 10)

x = torch.randn(10, 10)
z = torch.randn(10, 10)

mask1 = torch.randint(0, 2,(10, 10)).float()
mask2 = torch.randint(0, 2,(10, 10)).float()
# mask2 = torc…