Replacing a block diagonal with other values in a tensor

Hi everyone,

I am wondering how to replace a block diagonal matrix with other values in a tensor in a batched manner. For example, given a tensor with shape [B, 3, 6], where B is the batch size (here B = 1 for illustration):
[[[ 0, 1, 2, 3, 4, 5],
[ 6, 7, 8, 9, 10, 11],
[12, 13, 14, 15, 16, 17]]]

Now I would like to replace each 1 x 2 block along the diagonal with a new value, say [1,1], which transforms the tensor above to
[[[ 1, 1, 2, 3, 4, 5],
[ 6, 7, 1, 1, 10, 11],
[12, 13, 14, 15, 1, 1]]]

Is it a must to achieve this through a loop or I just ignored some efficient APIs in PyTorch?

Thank you very much!