Hey guys,
I am currently working on converting a Tensorflow project to PyTorch. At the end of the prediction, I get a translation vector field and rotation vector field of the sizes [B, 3, h, w] and [B, 3, 3, h, w] respectively.
In the Tensorflow version, the outputs are [B, h, w, 3] for the translation field and [B, h, w, 3, 3] for the rotation field.
I now need to perform a matrix multiplication on two translation fields. The Tensorflow code would look like this:
# with rot_mat1 and rot_mat being of size **[B, h, w, 3, 3]**
r2r1 = tf.matmul(rot_mat2, rot_mat1)      

All working fine.
When I do the same without thinking in PyTorch:
# with rot_mat1 and rot_mat being of size **[B, 3, 3, h, w]**
r2r1 = tf.matmul(rot_mat2, rot_mat1)  

I get the error:
File "/home/student/lukas/pycharm-sync/depth-from-video/utils.py", line 131, in combine
    r2r1 = torch.matmul(rot_mat2, rot_mat1)
RuntimeError: invalid argument 6: wrong matrix size at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:558

I know the problem is the different dimensions in PyTorch. But do you guys have any quick fix for how I get the desired result as in the Tensorflow version? I know the answer is most likely trivial, but I am somehow stuck at the moment.
Thanks so much!