Как использовать функцию torch.stack

Question

У меня есть вопрос о torch.stack

У меня есть 2 тензора, a.shape = (2, 3, 4) и b.shape = (2, 3). Как их складывать без операции на месте?

Answer 1

Для укладки требуется одинаковое количество измерений.Одним из способов было бы отжать и сложить.Например:

a.size()  # 2, 3, 4
b.size()  # 2, 3
b = torch.unsqueeze(b, dim=2)  # 2, 3, 1
# torch.unsqueeze(b, dim=-1) does the same thing

torch.stack([a, b], dim=2)  # 2, 3, 5