Possible Miscalculation of MSELoss in Face Generator

[cchadj]

I would like to report a possible miscalculation of the loss in the face generator.

Issue description

Please have a look at the following code snippet:

TalkSHOW/nets/smplx_face.py

Lines 155 to 159 in 38aab30

	MSELoss = torch.mean(torch.abs(pred_poses[:, :, :6] - gt_poses[:, :, :6]))
	if self.expression:
	expl = torch.mean((pred_poses[:, :, -100:] - gt_poses[:, :, -100:])**2)
	else:
	expl = 0

I believe the loss calculation at line 155 is wrong. The slice should go up to index 3, not 6.
That's because the dimensions for the jaw_pose are 3.

I would like to remind you that pred_poses shape is (N, seq_length, 103), where the first 3 dimensions are for the jaw_pose while the rest 100 are for the expression.

For the gt_poses the shape is (N, seq_length, 265) where the first 3 dimensions are for the jaw pose and the last 100 are for expression.
The 3 next dimensions after the first 3 of the jaw pose are for the left eye.

When we do MSELoss = torch.mean(torch.abs(pred_poses[:, :, :6] - gt_poses[:, :, :6])) we compare correctly the first 3 jaw_pose features but also we compare 3 left eye features from gt_poses with 3 features expression features from pred_poses.

Proposed Fix:

I believe the correct way to calculate the loss is by changing 6 to 3, as follows:
MSELoss = torch.mean(torch.abs(pred_poses[:, :, :3] - gt_poses[:, :, :3])).

Please let me know if my assertion is correct or whether I misunderstood something.