diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml
new file mode 100644
index 0000000..db86793
--- /dev/null
+++ b/.github/workflows/coverage.yml
@@ -0,0 +1,38 @@
+name: Run tests and upload coverage
+
+on: 
+  [workflow_dispatch, pull_request]
+
+jobs:
+  test:
+    name: Run tests and collect coverage
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 2
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+
+      - name: Install repo dependencies
+        run: pip install .
+
+      - name: Install test dependencies
+        run: pip install pytest pytest-cov
+
+      - name: Test with pytest
+        run: |
+          pytest --cov --junitxml=junit.xml -o junit_family=legacy
+
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v5
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+
+      - name: Upload test results to Codecov
+        if: ${{ !cancelled() }}
+        uses: codecov/test-results-action@v1
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
\ No newline at end of file
diff --git a/README.md b/README.md
index 130835c..3991ff8 100644
--- a/README.md
+++ b/README.md
@@ -28,7 +28,7 @@ You can also clone this repository and run the notebooks locally with [Jupyter](
 
 ## Quickstart
 
-```
+```python
 import torch
 
 from lucent.optvis import render
diff --git a/demo.py b/demo.py
index 380e09d..f7b2aa7 100644
--- a/demo.py
+++ b/demo.py
@@ -16,12 +16,12 @@ def main():
 
     if CPPN:
         # CPPN parameterization
-        param_f = lambda: param.cppn(224)
+        param_f = lambda: param.cppn(224, device=device)
         opt = lambda params: torch.optim.Adam(params, 5e-3)
         # Some objectives work better with CPPN than others
         obj = "mixed4d_3x3_bottleneck_pre_relu_conv:139"
     else:
-        param_f = lambda: param.image(224, fft=SPATIAL_DECORRELATION, decorrelate=CHANNEL_DECORRELATION)
+        param_f = lambda: param.image(224, fft=SPATIAL_DECORRELATION, decorrelate=CHANNEL_DECORRELATION, device=device)
         opt = lambda params: torch.optim.Adam(params, 5e-2)
         obj = "mixed4a:476"
 
diff --git a/lucent/optvis/objectives.py b/lucent/optvis/objectives.py
index 2233359..b7a79d8 100644
--- a/lucent/optvis/objectives.py
+++ b/lucent/optvis/objectives.py
@@ -231,14 +231,13 @@ def inner(model):
 
 
 def _torch_blur(tensor, out_c=3):
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     depth = tensor.shape[1]
     weight = np.zeros([depth, depth, out_c, out_c])
     for ch in range(depth):
         weight_ch = weight[ch, ch, :, :]
         weight_ch[ :  ,  :  ] = 0.5
         weight_ch[1:-1, 1:-1] = 1.0
-    weight_t = torch.tensor(weight).float().to(device)
+    weight_t = torch.tensor(weight).float().to(tensor.device)
     conv_f = lambda t: F.conv2d(t, weight_t, None, 1, 1)
     return conv_f(tensor) / conv_f(torch.ones_like(tensor))
 
diff --git a/lucent/optvis/param/color.py b/lucent/optvis/param/color.py
index deab013..94d449d 100644
--- a/lucent/optvis/param/color.py
+++ b/lucent/optvis/param/color.py
@@ -31,9 +31,8 @@
 
 
 def _linear_decorrelate_color(tensor):
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     t_permute = tensor.permute(0, 2, 3, 1)
-    t_permute = torch.matmul(t_permute, torch.tensor(color_correlation_normalized.T).to(device))
+    t_permute = torch.matmul(t_permute, torch.tensor(color_correlation_normalized.T).to(t_permute.device))
     tensor = t_permute.permute(0, 3, 1, 2)
     return tensor
 
diff --git a/lucent/optvis/param/cppn.py b/lucent/optvis/param/cppn.py
index 65b08cd..82622ad 100644
--- a/lucent/optvis/param/cppn.py
+++ b/lucent/optvis/param/cppn.py
@@ -19,6 +19,8 @@
 import torch
 import numpy as np
 
+from lucent.util import DEFAULT_DEVICE
+
 
 class CompositeActivation(torch.nn.Module):
 
@@ -28,7 +30,8 @@ def forward(self, x):
 
 
 def cppn(size, num_output_channels=3, num_hidden_channels=24, num_layers=8,
-         activation_fn=CompositeActivation, normalize=False):
+         activation_fn=CompositeActivation, normalize=False,
+         device=DEFAULT_DEVICE):
 
     r = 3 ** 0.5
 
@@ -36,7 +39,7 @@ def cppn(size, num_output_channels=3, num_hidden_channels=24, num_layers=8,
     x = coord_range.view(-1, 1).repeat(1, coord_range.size(0))
     y = coord_range.view(1, -1).repeat(coord_range.size(0), 1)
 
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    device = torch.device(device)
 
     input_tensor = torch.stack([x, y], dim=0).unsqueeze(0).to(device)
 
diff --git a/lucent/optvis/param/gan.py b/lucent/optvis/param/gan.py
index c181172..fb684c0 100644
--- a/lucent/optvis/param/gan.py
+++ b/lucent/optvis/param/gan.py
@@ -27,6 +27,24 @@
             "fc7": "https://onedrive.live.com/download?cid=9CFFF6BCB39F6829&resid=9CFFF6BCB39F6829%2145338&authkey=AJ0R-daUAVYjQIw",
             "fc8": "https://onedrive.live.com/download?cid=9CFFF6BCB39F6829&resid=9CFFF6BCB39F6829%2145340&authkey=AKIfNk7s5MGrRkU"}
 
+def download_url_to_file_fake_request(url, dst):
+    """
+    Download object at the given URL to a local path, using browser-like HTTP GET request.
+    """
+
+    import requests
+    from tqdm import tqdm
+
+    # Imitate Chrome browser
+    headers = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_14_6) "
+                                "AppleWebKit/537.36 (KHTML, like Gecko) "
+                                "Chrome/76.0.3809.132 Safari/537.36"}
+
+    with requests.get(url, headers=headers, stream=True) as r:
+        r.raise_for_status()
+        with open(dst, 'wb') as f:
+            for chunk in tqdm(r.iter_content(chunk_size=8192)):
+                f.write(chunk)
 
 def load_statedict_from_online(name="fc6"):
     torchhome = torch.hub._get_torch_home()
@@ -34,8 +52,12 @@ def load_statedict_from_online(name="fc6"):
     os.makedirs(ckpthome, exist_ok=True)
     filepath = join(ckpthome, "upconvGAN_%s.pt"%name)
     if not os.path.exists(filepath):
-        torch.hub.download_url_to_file(model_urls[name], filepath, hash_prefix=None,
-                                   progress=True)
+        print("Downloading %s"%model_urls[name])
+        download_url_to_file_fake_request(model_urls[name], filepath)
+
+        # this is blocked by onedrive
+        #torch.hub.download_url_to_file(model_urls[name], filepath, hash_prefix=None,
+        #                           progress=True)
     SD = torch.load(filepath)
     return SD
 
diff --git a/lucent/optvis/param/images.py b/lucent/optvis/param/images.py
index b11b666..5385b2a 100644
--- a/lucent/optvis/param/images.py
+++ b/lucent/optvis/param/images.py
@@ -19,16 +19,17 @@
 
 from lucent.optvis.param.spatial import pixel_image, fft_image
 from lucent.optvis.param.color import to_valid_rgb
+from lucent.util import DEFAULT_DEVICE
 
 
 def image(w, h=None, sd=None, batch=None, decorrelate=True,
-          fft=True, channels=None):
+          fft=True, channels=None, device=DEFAULT_DEVICE):
     h = h or w
     batch = batch or 1
     ch = channels or 3
     shape = [batch, ch, h, w]
     param_f = fft_image if fft else pixel_image
-    params, image_f = param_f(shape, sd=sd)
+    params, image_f = param_f(shape, sd=sd, device=device)
     if channels:
         output = to_valid_rgb(image_f, decorrelate=False)
     else:
diff --git a/lucent/optvis/param/lowres.py b/lucent/optvis/param/lowres.py
index e5aed0d..04aed64 100644
--- a/lucent/optvis/param/lowres.py
+++ b/lucent/optvis/param/lowres.py
@@ -22,9 +22,11 @@
 import torch.nn.functional as F
 
 from lucent.optvis.param.resize_bilinear_nd import resize_bilinear_nd
+from lucent.util import DEFAULT_DEVICE
 
 
-def lowres_tensor(shape, underlying_shape, offset=None, sd=0.01):
+def lowres_tensor(shape, underlying_shape, offset=None, sd=0.01,
+                  device=DEFAULT_DEVICE):
     """Produces a tensor paramaterized by a interpolated lower resolution tensor.
     This is like what is done in a laplacian pyramid, but a bit more general. It
     can be a powerful way to describe images.
@@ -41,7 +43,7 @@ def lowres_tensor(shape, underlying_shape, offset=None, sd=0.01):
     Returns:
         A tensor paramaterized by a lower resolution tensorflow variable.
     """
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    device = torch.device(device)
     underlying_t = (torch.randn(*underlying_shape) * sd).to(device).requires_grad_(True)
     if offset is not None:
         # Deal with non-list offset
diff --git a/lucent/optvis/param/spatial.py b/lucent/optvis/param/spatial.py
index 8fd4460..414c9c2 100644
--- a/lucent/optvis/param/spatial.py
+++ b/lucent/optvis/param/spatial.py
@@ -18,12 +18,13 @@
 import torch
 import numpy as np
 
+from lucent.util import DEFAULT_DEVICE
+
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 TORCH_VERSION = torch.__version__
 
 
-def pixel_image(shape, sd=None):
+def pixel_image(shape, sd=None, device=DEFAULT_DEVICE):
     sd = sd or 0.01
     tensor = (torch.randn(*shape) * sd).to(device).requires_grad_(True)
     return [tensor], lambda: tensor
@@ -42,7 +43,7 @@ def rfft2d_freqs(h, w):
     return np.sqrt(fx * fx + fy * fy)
 
 
-def fft_image(shape, sd=None, decay_power=1):
+def fft_image(shape, sd=None, decay_power=1, device=DEFAULT_DEVICE):
     batch, channels, h, w = shape
     freqs = rfft2d_freqs(h, w)
     init_val_size = (batch, channels) + freqs.shape + (2,) # 2 for imaginary and real components
diff --git a/lucent/optvis/render.py b/lucent/optvis/render.py
index 4c15168..ec04723 100644
--- a/lucent/optvis/render.py
+++ b/lucent/optvis/render.py
@@ -43,7 +43,7 @@ def render_vis(
     fixed_image_size=None,
 ):
     if param_f is None:
-        param_f = lambda: param.image(128)
+        param_f = lambda: param.image(128, device=next(model.parameters()).device)
     # param_f is a function that should return two things
     # params - parameters to update, which we pass to the optimizer
     # image_f - a function that returns an image as a tensor
@@ -81,7 +81,7 @@ def render_vis(
 
     transform_f = transform.compose(transforms)
 
-    hook = hook_model(model, image_f)
+    hook, features = hook_model(model, image_f, return_hooks=True)
     objective_f = objectives.as_objective(objective_f)
 
     if verbose:
@@ -124,6 +124,10 @@ def closure():
             print("Loss at step {}: {:.3f}".format(i, objective_f(hook)))
         images.append(tensor_to_img_array(image_f()))
 
+    # Clear hooks
+    for module_hook in features.values():
+        del module_hook.module._forward_hooks[module_hook.hook.id]
+
     if save_image:
         export(image_f(), image_name)
     if show_inline:
@@ -136,6 +140,9 @@ def closure():
 def tensor_to_img_array(tensor):
     image = tensor.cpu().detach().numpy()
     image = np.transpose(image, [0, 2, 3, 1])
+    # Check if the image is single channel and convert to 3-channel
+    if len(image.shape) == 4 and image.shape[3] == 1:  # Single channel image
+        image = np.repeat(image, 3, axis=3)
     return image
 
 
@@ -149,6 +156,8 @@ def view(tensor):
     image = (image * 255).astype(np.uint8)
     if len(image.shape) == 4:
         image = np.concatenate(image, axis=1)
+    if len(image.shape) == 3 and image.shape[2] == 1:
+        image = image.squeeze(2)
     Image.fromarray(image).show()
 
 
@@ -177,10 +186,11 @@ def hook_fn(self, module, input, output):
         self.features = output
 
     def close(self):
+        # This doesn't actually do anything
         self.hook.remove()
 
 
-def hook_model(model, image_f):
+def hook_model(model, image_f, return_hooks=False):
     features = OrderedDict()
 
     # recursive hooking function
@@ -206,4 +216,6 @@ def hook(layer):
         assert out is not None, "There are no saved feature maps. Make sure to put the model in eval mode, like so: `model.to(device).eval()`. See README for example."
         return out
 
+    if return_hooks:
+        return hook, features
     return hook
diff --git a/lucent/optvis/transform.py b/lucent/optvis/transform.py
index ad8aab3..0ddba33 100644
--- a/lucent/optvis/transform.py
+++ b/lucent/optvis/transform.py
@@ -22,8 +22,13 @@
 import kornia
 from kornia.geometry.transform import translate
 
+try:
+    from kornia import warp_affine, get_rotation_matrix2d
+except ImportError:
+    from kornia.geometry.transform import warp_affine, get_rotation_matrix2d
+
+
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 KORNIA_VERSION = kornia.__version__
 
 
@@ -33,7 +38,7 @@ def jitter(d):
     def inner(image_t):
         dx = np.random.choice(d)
         dy = np.random.choice(d)
-        return translate(image_t, torch.tensor([[dx, dy]]).float().to(device))
+        return translate(image_t, torch.tensor([[dx, dy]]).float().to(image_t.device))
 
     return inner
 
@@ -52,7 +57,7 @@ def random_scale(scales):
     def inner(image_t):
         scale = np.random.choice(scales)
         shp = image_t.shape[2:]
-        scale_shape = [_roundup(scale * d) for d in shp]
+        scale_shape = [int(_roundup(scale * d)) for d in shp]
         pad_x = max(0, _roundup((shp[1] - scale_shape[1]) / 2))
         pad_y = max(0, _roundup((shp[0] - scale_shape[0]) / 2))
         upsample = torch.nn.Upsample(
@@ -76,8 +81,8 @@ def inner(image_t):
         center = torch.ones(b, 2)
         center[..., 0] = (image_t.shape[3] - 1) / 2
         center[..., 1] = (image_t.shape[2] - 1) / 2
-        M = kornia.get_rotation_matrix2d(center, angle, scale).to(device)
-        rotated_image = kornia.warp_affine(image_t.float(), M, dsize=(h, w))
+        M = get_rotation_matrix2d(center, angle, scale).to(image_t.device)
+        rotated_image = warp_affine(image_t.float(), M, dsize=(h, w))
         return rotated_image
 
     return inner
diff --git a/lucent/util.py b/lucent/util.py
index 2ce264c..ac0f40b 100644
--- a/lucent/util.py
+++ b/lucent/util.py
@@ -17,13 +17,18 @@
 
 from __future__ import absolute_import, division, print_function
 
+import numpy as np
 import torch
 import random
 
 
+DEFAULT_DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
 def set_seed(seed):
     # Set global seeds to for reproducibility
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
     torch.backends.cudnn.deterministic=True
     random.seed(seed)
+    np.random.seed(seed)
diff --git a/setup.py b/setup.py
index 597d0eb..8ec8692 100644
--- a/setup.py
+++ b/setup.py
@@ -33,7 +33,7 @@
     install_requires=[
         "torch>=1.5.0",
         "torchvision",
-        "kornia<=0.4.1",
+        "kornia>=0.4.1",
         "tqdm",
         "numpy",
         "ipython",
diff --git a/tests/optvis/param/test_cppn.py b/tests/optvis/param/test_cppn.py
index 8a86808..08a8899 100644
--- a/tests/optvis/param/test_cppn.py
+++ b/tests/optvis/param/test_cppn.py
@@ -15,11 +15,8 @@
 
 from __future__ import absolute_import, division, print_function
 
-import pytest
-
 import torch
-import numpy as np
-from lucent.optvis import param, render, objectives
+from lucent.optvis import param, objectives
 
 
 def xor_loss(T):
diff --git a/tests/optvis/param/test_gan.py b/tests/optvis/param/test_gan.py
index 4ee7747..af0e17a 100644
--- a/tests/optvis/param/test_gan.py
+++ b/tests/optvis/param/test_gan.py
@@ -1,28 +1,28 @@
 import pytest
 
 import torch
-import numpy as np
-from lucent.optvis import param, render, objectives
+from lucent.optvis import render, objectives
 from lucent.optvis.param.gan import upconvGAN
 from lucent.modelzoo import inceptionv1
+from lucent.util import DEFAULT_DEVICE
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 NUM_STEPS = 5
 
+
 @pytest.fixture
 def inceptionv1_model():
-    model = inceptionv1().to(device).eval()
+    model = inceptionv1().to(DEFAULT_DEVICE).eval()
     return model
 
 def test_fc6gan_load():
     """ Test if gan could be downloaded and loaded
     It will download the model and store it locally
     """
-    G = upconvGAN("fc6").to(device)
+    G = upconvGAN("fc6").to(DEFAULT_DEVICE)
 
     def GANparam(batch=1, sd=1):
-        code = (torch.randn((batch, G.codelen)) * sd).to(device).requires_grad_(True)
+        code = (torch.randn((batch, G.codelen)) * sd).to(DEFAULT_DEVICE).requires_grad_(True)
         imagef = lambda:  G.visualize(code)
         return [code], imagef
     code, imagef = GANparam(batch=2, sd=1)
@@ -30,10 +30,10 @@ def GANparam(batch=1, sd=1):
     assert img.shape == (2, 3, 256, 256), "Cannot forward fc6 GAN, shape incorrect."
 
 def test_fc7gan_load():
-    G = upconvGAN("fc7").to(device)
+    G = upconvGAN("fc7").to(DEFAULT_DEVICE)
 
     def GANparam(batch=1, sd=1):
-        code = (torch.randn((batch, G.codelen)) * sd).to(device).requires_grad_(True)
+        code = (torch.randn((batch, G.codelen)) * sd).to(DEFAULT_DEVICE).requires_grad_(True)
         imagef = lambda: G.visualize(code)
         return [code], imagef
     code, imagef = GANparam(batch=2, sd=1)
@@ -41,10 +41,10 @@ def GANparam(batch=1, sd=1):
     assert img.shape == (2, 3, 256, 256), "Cannot forward fc7 GAN, shape incorrect."
 
 def test_fc8gan_load():
-    G = upconvGAN("fc8").to(device)
+    G = upconvGAN("fc8").to(DEFAULT_DEVICE)
 
     def GANparam(batch=1, sd=1):
-        code = (torch.randn((batch, G.codelen)) * sd).to(device).requires_grad_(True)
+        code = (torch.randn((batch, G.codelen)) * sd).to(DEFAULT_DEVICE).requires_grad_(True)
         imagef = lambda: G.visualize(code)
         return [code], imagef
     code, imagef = GANparam(batch=2, sd=1)
@@ -52,10 +52,10 @@ def GANparam(batch=1, sd=1):
     assert img.shape == (2, 3, 256, 256), "Cannot forward fc8 GAN, shape incorrect."
 
 def test_pool5gan_load():
-    G = upconvGAN("pool5").to(device)
+    G = upconvGAN("pool5").to(DEFAULT_DEVICE)
 
     def GANparam(batch=1, sd=1):
-        code = (torch.randn((batch, G.codelen, 6, 6)) * sd).to(device).requires_grad_(True)
+        code = (torch.randn((batch, G.codelen, 6, 6)) * sd).to(DEFAULT_DEVICE).requires_grad_(True)
         imagef = lambda: G.visualize(code)
         return [code], imagef
     code, imagef = GANparam(batch=2, sd=1)
@@ -82,10 +82,10 @@ def assert_gan_gradient_descent(GANparam, objective, model):
 
 def test_gan_img_optim(inceptionv1_model):
     """ Test if GAN generated image could be optimized """
-    G = upconvGAN("fc6").to(device)
+    G = upconvGAN("fc6").to(DEFAULT_DEVICE)
 
     def GANparam(batch=1, sd=1):
-        code = (torch.randn((batch, G.codelen)) * sd).to(device).requires_grad_(True)
+        code = (torch.randn((batch, G.codelen)) * sd).to(DEFAULT_DEVICE).requires_grad_(True)
         imagef = lambda:  G.visualize(code)
         return [code], imagef
     objective = objectives.neuron("input", 0)
@@ -94,10 +94,10 @@ def GANparam(batch=1, sd=1):
 
 def test_gan_deep_optim(inceptionv1_model):
     """ Test if GAN generated image could be optimized """
-    G = upconvGAN("fc6").to(device)
+    G = upconvGAN("fc6").to(DEFAULT_DEVICE)
 
     def GANparam(batch=1, sd=1):
-        code = (torch.randn((batch, G.codelen)) * sd).to(device).requires_grad_(True)
+        code = (torch.randn((batch, G.codelen)) * sd).to(DEFAULT_DEVICE).requires_grad_(True)
         imagef = lambda:  G.visualize(code)
         return [code], imagef
 
diff --git a/tests/optvis/param/test_lowres.py b/tests/optvis/param/test_lowres.py
index 06410f3..697cf92 100644
--- a/tests/optvis/param/test_lowres.py
+++ b/tests/optvis/param/test_lowres.py
@@ -15,8 +15,6 @@
 
 from __future__ import absolute_import, division, print_function
 
-import pytest
-
 from lucent.optvis import param
 
 
diff --git a/tests/optvis/test_integration.py b/tests/optvis/test_integration.py
index fb60b8c..63c24ce 100644
--- a/tests/optvis/test_integration.py
+++ b/tests/optvis/test_integration.py
@@ -18,13 +18,14 @@
 import pytest
 
 import torch
-from lucent.optvis import objectives, param, render, transform
+from lucent.optvis import param, render
 from lucent.modelzoo import inceptionv1
+from lucent.util import DEFAULT_DEVICE
 
 
 @pytest.fixture
 def inceptionv1_model():
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    device = torch.device(DEFAULT_DEVICE)
     model = inceptionv1().to(device).eval()
     return model
 
diff --git a/tests/optvis/test_objectives.py b/tests/optvis/test_objectives.py
index f27a72e..45d69ee 100644
--- a/tests/optvis/test_objectives.py
+++ b/tests/optvis/test_objectives.py
@@ -18,18 +18,17 @@
 import pytest
 
 import torch
-import random
-import numpy as np
 from lucent.util import set_seed
 from lucent.optvis import objectives, param, render
 from lucent.modelzoo import inceptionv1
+from lucent.util import DEFAULT_DEVICE
 
 
 set_seed(137)
 
 
 NUM_STEPS = 5
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+device = torch.device(DEFAULT_DEVICE)
 
 @pytest.fixture
 def inceptionv1_model():
@@ -120,12 +119,12 @@ def test_diversity(inceptionv1_model):
 
 
 def test_direction(inceptionv1_model):
-    direction = torch.rand(512) * 1000
+    direction = torch.rand(512, device=next(inceptionv1_model.parameters()).device) * 1000
     objective = objectives.direction(layer='mixed4c', direction=direction)
     assert_gradient_descent(objective, inceptionv1_model)
 
 
 def test_direction_neuron(inceptionv1_model):
-    direction = torch.rand(512) * 1000
+    direction = torch.rand(512, device=next(inceptionv1_model.parameters()).device) * 1000
     objective = objectives.direction_neuron(layer='mixed4c', direction=direction)
     assert_gradient_descent(objective, inceptionv1_model)
diff --git a/tests/optvis/test_render.py b/tests/optvis/test_render.py
index 77392ea..a9a7d3f 100644
--- a/tests/optvis/test_render.py
+++ b/tests/optvis/test_render.py
@@ -17,15 +17,14 @@
 
 import pytest
 
-import torch
 from lucent.optvis import render, param
 from lucent.modelzoo import inceptionv1
+from lucent.util import DEFAULT_DEVICE
 
 
 @pytest.fixture
 def inceptionv1_model():
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    model = inceptionv1().to(device).eval()
+    model = inceptionv1().to(DEFAULT_DEVICE).eval()
     return model
 
 
diff --git a/tests/optvis/test_transform.py b/tests/optvis/test_transform.py
index c1ce069..8741711 100644
--- a/tests/optvis/test_transform.py
+++ b/tests/optvis/test_transform.py
@@ -15,46 +15,42 @@
 
 from __future__ import absolute_import, division, print_function
 
-import pytest
-
 import torch
 import numpy as np
 from lucent.optvis import transform
-
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+from lucent.util import DEFAULT_DEVICE
 
 
 def test_pad_reflect():
     pad = transform.pad(1)
-    tensor = torch.ones(1, 3, 2, 2).to(device)
-    assert torch.all(pad(tensor).eq(torch.ones(1, 3, 4, 4).to(device)))
+    tensor = torch.ones(1, 3, 2, 2).to(DEFAULT_DEVICE)
+    assert torch.all(pad(tensor).eq(torch.ones(1, 3, 4, 4).to(DEFAULT_DEVICE)))
 
 
 def test_pad_constant():
     pad = transform.pad(1, mode="constant")
-    tensor = torch.ones(1, 3, 2, 2).to(device)
+    tensor = torch.ones(1, 3, 2, 2).to(DEFAULT_DEVICE)
     assert torch.all(pad(tensor).eq(torch.tensor([[
         [[0.5, 0.5, 0.5, 0.5], [0.5, 1, 1, 0.5], [0.5, 1, 1, 0.5], [0.5, 0.5, 0.5, 0.5]],
         [[0.5, 0.5, 0.5, 0.5], [0.5, 1, 1, 0.5], [0.5, 1, 1, 0.5], [0.5, 0.5, 0.5, 0.5]],
         [[0.5, 0.5, 0.5, 0.5], [0.5, 1, 1, 0.5], [0.5, 1, 1, 0.5], [0.5, 0.5, 0.5, 0.5]],
-    ]]).to(device)))
+    ]]).to(DEFAULT_DEVICE)))
 
 
 def test_random_scale_down():
     scale = transform.random_scale([0.33])
-    tensor = torch.ones(1, 3, 3, 3).to(device)
+    tensor = torch.ones(1, 3, 3, 3).to(DEFAULT_DEVICE)
     assert torch.all(scale(tensor).eq(torch.tensor([[
         [[0, 0, 0], [0, 1, 0], [0, 0, 0]],
         [[0, 0, 0], [0, 1, 0], [0, 0, 0]],
         [[0, 0, 0], [0, 1, 0], [0, 0, 0]]
-    ]]).to(device)))
+    ]]).to(DEFAULT_DEVICE)))
 
 
 def test_random_scale_up():
     scale = transform.random_scale([2])
-    tensor = torch.ones(1, 3, 1, 1).to(device)
-    assert torch.all(scale(tensor).eq(torch.ones(1, 3, 2, 2).to(device)))
+    tensor = torch.ones(1, 3, 1, 1).to(DEFAULT_DEVICE)
+    assert torch.all(scale(tensor).eq(torch.ones(1, 3, 2, 2).to(DEFAULT_DEVICE)))
 
 
 def test_random_rotate_even_size():
@@ -63,12 +59,12 @@ def test_random_rotate_even_size():
         [[0, 1], [0, 1]],
         [[0, 1], [0, 1]],
         [[0, 1], [0, 1]],
-    ]]).to(device)
+    ]]).to(DEFAULT_DEVICE)
     assert torch.all(rotate(tensor).eq(torch.tensor([[
         [[1, 1], [0, 0]],
         [[1, 1], [0, 0]],
         [[1, 1], [0, 0]],
-    ]]).to(device)))
+    ]]).to(DEFAULT_DEVICE)))
 
 
 def test_random_rotate_odd_size():
@@ -77,20 +73,20 @@ def test_random_rotate_odd_size():
         [[0, 0, 1], [0, 0, 1], [0, 0, 1]],
         [[0, 0, 1], [0, 0, 1], [0, 0, 1]],
         [[0, 0, 1], [0, 0, 1], [0, 0, 1]]
-    ]]).to(device)
+    ]]).to(DEFAULT_DEVICE)
     assert torch.all(rotate(tensor).eq(torch.tensor([[
         [[1, 1, 1], [0, 0, 0], [0, 0, 0]],
         [[1, 1, 1], [0, 0, 0], [0, 0, 0]],
         [[1, 1, 1], [0, 0, 0], [0, 0, 0]]
-    ]]).to(device)))
+    ]]).to(DEFAULT_DEVICE)))
 
 
 def test_normalize():
     normalize = transform.normalize()
-    tensor = torch.zeros(1, 3, 1, 1).to(device)
+    tensor = torch.zeros(1, 3, 1, 1).to(DEFAULT_DEVICE)
     print(normalize(tensor))
     assert torch.allclose(normalize(tensor), torch.tensor([[
         [[-0.485/0.229]],
         [[-0.456/0.224]],
         [[-0.406/0.225]]
-    ]]).to(device))
+    ]]).to(DEFAULT_DEVICE))