conv type为2d时求theta是不是有点问题

[WuChannn]

原代码中：
if self.conv_type == '2d':
theta = torch.mean(theta, axis=-1)
但是这样的话theta为三维的，是不是应该在求mean之前reshape，如：
if self.conv_type == '2d':
theta = torch.reshape(theta, (:,:,-1))
theta = torch.mean(theta, axis=-1)

[WuChannn]

@bobo0810 期待与您探讨一下

[bobo0810]

@bobo0810 期待与您探讨一下

你好，非常抱歉，我在训练过程中没有遇到这个问题。也许你可以去官方库探讨一下？