You should be able to use diagonal for assignment:
a = torch.zeros(6, 5, 5, 7) 
b = torch.randn(6, 5, 7) 
# the diagonal dimension is always last for numpy compat...
d = a.diagonal(dim1=1, dim2=2).permute(0, 2, 1)
d.copy_(b)
# now the "diagonal" of a is set to b.

Best regards
Thomas