Solved by ptrblck in post #19 
              
              
                Thanks again for the code snippet as well as the great debugging! 
I was able to narrow it down to a sync issue using this “minimal” code snippet: 
x = torch.randn(1, 2, 4, 4, device='cuda')
v = 2
m = torch.randn(1024, 1024, device='cuda')

res = []
for _ in range(10):
    psv_tgt = x[:, 0:1].repeat…