LearnerYme:

I guess that’s because torch.cat() would break the gradient,


That’s not the case as seen in this minimal code snippet:
x1 = torch.randn(2, 3, requires_grad=True)
x2 = torch.randn(2, 1, requires_grad=True)

y = torch.cat((x1, x2), dim=1)
y.mean().backward()

print(x1.grad)
print(x2.grad)

To verify it, you could also check the .grad attributes of your model and check, if another operation breaks the computation graph.