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!