Transform3d Stacking

[davezdeng8]

Unless I am understanding stacking incorrectly, there seems to be an issue with stacking Transform3D.

To reproduce:

import torch
from pytorch3d.transforms import Transform3D
transform3 = Transform3d().rotate(torch.stack([torch.eye(3)]*3)).translate(torch.zeros(3,3))
transform1 = Transform3d()
transform4 = transform1.stack(transform3)
print(len(transform3))
print(len(transform1))
print(len(transform4))
transform4.transform_points(torch.zeros(4,5,3))

The last line errors with:

ValueError: Expected batch dim for bmm to be equal or 1; got torch.Size([4, 5, 4]), torch.Size([2, 4, 4])

[bottler]

This is a bug. The stack function isn't doing the right thing when one of its inputs is a composition (even if the batch sizes match). You can get round this by replacing

transform4 = transform1.stack(transform3)

with

transform4 = transform1.stack(Transform3d(matrix=transform3.get_matrix()))