vae.py: throwing AssertionError at training time

[wisclmy0611]

Hi Choy. Thanks for your great work on this project! I'm trying to train the model in vae.py. I ran python -m examples.vae --train but got the following AssertionError:

Traceback (most recent call last):
  File "/path/to/anaconda3/envs/py3-mink/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/path/to/anaconda3/envs/py3-mink/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/path/to/MinkowskiEngine/examples/vae.py", line 737, in <module>
    train(net, dataloader, device, config)
  File "/path/to/MinkowskiEngine/examples/vae.py", line 632, in train
    out_cls, targets, sout, means, log_vars, zs = net(sin, target_key)
  File "/path/to/anaconda3/envs/py3-mink/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/path/to/MinkowskiEngine/examples/vae.py", line 588, in forward
    zs += torch.exp(0.5 * log_vars.F) * torch.randn_like(log_vars.F)
  File "/path/to/MinkowskiEngine/SparseTensor.py", line 544, in __iadd__
    assert isinstance(other, SparseTensor)
AssertionError

MinkowskiEngine/examples/vae.py

Lines 584 to 590 in 2bb627e

	def forward(self, sinput, gt_target):
	means, log_vars = self.encoder(sinput)
	zs = means
	if self.training:
	zs += torch.exp(0.5 * log_vars.F) * torch.randn_like(log_vars.F)
	out_cls, targets, sout = self.decoder(zs, gt_target)
	return out_cls, targets, sout, means, log_vars, zs

At training time, line 588 will be executed

zs += torch.exp(0.5 * log_vars.F) * torch.randn_like(log_vars.F)

while at test time it won't, so there will be no problem during inference. This error is caused by a PyTorch Tensor added to a SparseTensor. Since you have released a pretrained checkpoint,

MinkowskiEngine/examples/vae.py

Lines 739 to 743 in 2bb627e

	if not os.path.exists(config.weights):
	logging.info(
	f'Downloaing pretrained weights. This might take a while...')
	urllib.request.urlretrieve(
	"https://bit.ly/39TvWys", filename=config.weights)

I assume you've successfully run the code to train a VAE. Could you please share how you have done that without encountering the error? Thanks so much!