somap is stored as a plain tensor in SOM.__init__ rather than being registered via register_buffer() or as an nn.Parameter. The same applies to coord.
This means:
model.to(device) won't move somap/coord automatically (the manual .to() override works around this but is fragile)
model.state_dict() won't include somap, so saving/loading a trained model loses the trained map
model.load_state_dict() can't restore the map
Suggested fix
Use self.register_buffer('somap', ...) for somap and self.register_buffer('coord', ...) for coord. This would also allow removing the manual .to() overrides in SOM, WSOM, and MCWSOM.
somapis stored as a plain tensor inSOM.__init__rather than being registered viaregister_buffer()or as annn.Parameter. The same applies tocoord.This means:
model.to(device)won't movesomap/coordautomatically (the manual.to()override works around this but is fragile)model.state_dict()won't includesomap, so saving/loading a trained model loses the trained mapmodel.load_state_dict()can't restore the mapSuggested fix
Use
self.register_buffer('somap', ...)forsomapandself.register_buffer('coord', ...)forcoord. This would also allow removing the manual.to()overrides inSOM,WSOM, andMCWSOM.