Solved by tom in post #2 
              
              
                How about 
mask = (A < torch.arange(0, A.size(1), device=A.device, dtype=A.dtype).unsqueeze(0))
masked_A = torch.where(mask, A, torch.zeros(1,1, dtype=A.dtype, device=A.device))

where 

the arange will give you a vector of j and unsqueezing makes is broadcasteable to As shape,
the comparison gives …