def test_calls_gpu_after_to_gpu(self):
device = torch.device("cuda")
mu = self.linear.mu
mu.to(device)
self.linear = noisy_linear.FactorizedNoisyLinear(mu)
self._test_calls(device)
def setUp(self, bias):
self.bias = bias
mu = torch.nn.Linear(6, 5, bias=self.bias)
self.linear = noisy_linear.FactorizedNoisyLinear(mu)