search

NVlabs / imaginaire

NVIDIA's Deep Imagination Team's PyTorch Library
Other
3.99k stars 444 forks source link

Model average - setting "self.dim = 0" when removing spectral norm #61

Closed schrute99 closed 2 years ago

schrute99 commented 3 years ago

Hi,

I have noticed that you set self.dim = 0 in the ModelAverage class (line 79).

This means that for the modules with spectral norm, the weights are not reshaped before computing the spectral norm (see line 175).

Shouldn't the dimension depend on the module type (linear, conv, etc) as it is here in line 264 ?

Thank you!

mingyuliutw commented 2 years ago

This is only applied to the module where the spectral norm is removed. It shouldn't be an issue.