VariationalDropout modules do not support CUDA device

[lcdavis13]

Reproduction code

import torch
from torchdeq.dropout import VariationalDropout1d

assert torch.cuda.is_available(), "CUDA is required for this repro."

device = torch.device("cuda:0")

drop = VariationalDropout1d(dropout=0.1).to(device)
drop.train()

x = torch.randn(2, 8, 16, device=device)  # CUDA input
print("input device:", x.device)

y = drop(x)  # This should work, but currently throws an exception

print("output device:", y.device)

Expected output

input device: cuda:0
output device: cuda:0

Actual output

input device: cuda:0
Traceback (most recent call last):
  File "{path}\dropout_failure_test.py", line 19, in <module>
    y = drop(x)
        ^^^^^^^
  File "{path2}\torch\nn\modules\module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "{path2}\torch\nn\modules\module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "{path2}\torchdeq\dropout.py", line 73, in forward
    return mask * x
           ~~~~~^~~
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Cause

In torchdeq/dropout.py many lines resemble the following:

m = torch.zeros(B, L, 1).bernoulli_(1 - self.dropout)

rather than specifying the device e.g.:

m = torch.zeros(B, L, 1, device=x.device).bernoulli_(1 - self.dropout)