you could also flatten the points to `[batch_size, 51]

So, I followed your advice to flatten the keypoints since I couldn’t figure out how to do the initial suggestion you could make sure the output and target have the shape [batch_size, 3, 17] (probably as misunderstanding on my part for how to do it).
I did it by setting self.kp_score = nn.Linear(66078720, 17*3*3), where 17 are the keypoints, 3 is number of images and the other 3 is to match the predicted keypoints shape [3, 51]. I flattened cat_kp = cat_kp.flatten(start_dim=0) and pred_kp = pred_kp.flatten(start_dim=0).
However, there are 2 “issues” that I am having now:

The kp_loss for is very high (over 200). Not sure if that’s an issue due to problem number 2, as I can’t begin to train the model.
I get an issue saying the GPU is out of memory. I have an RTX 2080 Ti 11GB. I thought this would be able to handle this. Is it because I am using conv3D, because it works fine to training for detection with a batch size of 8 images.