mask_setup = torch.ones(s, s) #Shape -> [s, s]

mask = torch.diag(mask_setup) #Shape -> [s*s, s*s]

mask = mask.reshape(s*s, s, s) #Shape -> [s*s, s, s]

This should give you what you need. There may be a simpler way to get it done though.