Great work!

[desertfex]

self.spa = self.one_hot(bs, num_joint, self.seg) #(bs, 25, 20) --->[bs, 20, 25, 25]
self.spa = self.spa.permute(0, 3, 2, 1).cuda() #[bs, 20, 25, 25]--->[bs, 25, 25, 20]

Can the second code be:
self.spa = self.spa.permute(0, 2, 3, 1).cuda() #[bs, 20, 25, 25]--->[bs, 25, 25, 20]

The dimensions are the same. Will this make any difference? Thank you.

[shuidongliu]

self.spa = self.one_hot(bs, num_joint, self.seg) #(bs, 25, 20) --->[bs, 20, 25, 25]
self.spa = self.spa.permute(0, 3, 2, 1).cuda() #[bs, 20, 25, 25]--->[bs, 25, 25, 20]

Can the second code be:
self.spa = self.spa.permute(0, 2, 3, 1).cuda() #[bs, 20, 25, 25]--->[bs, 25, 25, 20]

The dimensions are the same. Will this make any difference? Thank you.

You can use the pretrained model to test the modified code.