Not working when using GPU for miwae

[SiyiSun99]

Description

when using cuda with impute method "miwae", you will get error. However, it worked on my Mac with only cpu.

How to Reproduce

1. run your code with plugin = Imputers().get("miwae")
   2.See error

Expected Behavior

successfully run like other methods.

Error Context

Traceback (most recent call last):
File "/home/cs-sun2/CMIE_Project/Python_code/miwae_impute.py", line 77, in
out_train = plugin.fit_transform(train_miss.copy())
File "/home/cs-sun2/tabcsdi-env/lib/python3.8/site-packages/sklearn/utils/_set_output.py", line 157, in wrapped
data_to_wrap = f(self, X, args, kwargs)
File "/home/cs-sun2/tabcsdi-env/lib/python3.8/site-packages/sklearn/base.py", line 916, in fit_transform
return self.fit(X, fit_params).transform(X)
File "/home/cs-sun2/tabcsdi-env/lib/python3.8/site-packages/hyperimpute/plugins/core/base_plugin.py", line 124, in fit
return self._fit(X, args, kwargs)
File "/home/cs-sun2/tabcsdi-env/lib/python3.8/site-packages/hyperimpute/plugins/imputers/plugin_miwae.py", line 263, in _fit
loss = self._miwae_loss(iota_x=b_data, mask=b_mask)
File "/home/cs-sun2/tabcsdi-env/lib/python3.8/site-packages/hyperimpute/plugins/imputers/plugin_miwae.py", line 87, in miwaeloss
td.Normal(
File "/home/cs-sun2/tabcsdi-env/lib/python3.8/site-packages/torch/distributions/normal.py", line 57, in init
super().init(batch_shape, validate_args=validate_args)
File "/home/cs-sun2/tabcsdi-env/lib/python3.8/site-packages/torch/distributions/distribution.py", line 70, in init
raise ValueError(
ValueError: Expected parameter scale (Tensor of shape (256, 1)) of distribution Normal(loc: torch.Size([256, 1]), scale: torch.Size([256, 1])) to satisfy the constraint GreaterThan(lower_bound=0.0), but found invalid values:tensor([[1.1709e+00],
[1.1709e+00],
[1.4363e-12],
[0.0000e+00],
[0.0000e+00],
[0.0000e+00],
[1.1709e+00],
[0.0000e+00],
[0.0000e+00],
[0.0000e+00],
[0.0000e+00],
[1.1709e+00],
[0.0000e+00],
[0.0000e+00],
[1.1709e+00],
[1.1709e+00],
[0.0000e+00],
[1.1709e+00],
[1.1709e+00],
[1.1709e+00],
[0.0000e+00],
[0.0000e+00],
[0.0000e+00],
[1.1709e+00],
[1.1709e+00],
[1.1709e+00],
[1.1709e+00],
[1.1709e+00],
[0.0000e+00],
[1.1709e+00],
[0.0000e+00],
[0.0000e+00],
[1.1709e+00],
[1.1709e+00],
[1.1709e+00],
[1.1709e+00],
[0.0000e+00],
[0.0000e+00],
[1.1709e+00],
[0.0000e+00],
[0.0000e+00],
[0.0000e+00],
[1.1709e+00],
[0.0000e+00],
[1.1709e+00]], device='cuda:0', grad_fn=)

How to fix

in hyperimpute/plugins/imputers, plugin_miwae.py, in function _miwae_loss(..). when define variable q_zgivenxobs:

scale = torch.nn.Softplus()(
out_encoder[..., self.latent_size: (2*self.latent_size)]
)

however when you define the all_scales_obs_model:
all_scales_obs_model = torch.nn.Softplus()(
out_decoder[..., p:(2*p)]
) + 0.001.

You should also add 0.001 to q_zgivenxobs's scale parameter.