Fix broken PyTorch Lightning test

Closes #547 Ref #551 Signed-off-by: Ben Firshman <[email protected]>
replicate · Mar 8, 2021 · 95fd2b2 · 95fd2b2
1 parent 6bf33f2
commit 95fd2b2
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/python/tests/test_pl_callback.py b/python/tests/test_pl_callback.py
@@ -12,6 +12,7 @@
 from torchvision.datasets import MNIST
 from torchvision import transforms
 import pytest
+import urllib
 
 from keepsake.pl_callback import KeepsakeCallback
 
@@ -35,6 +36,12 @@ def forward(self, x):
         return x
 
     def prepare_data(self):
+        # HACK: https://github.com/pytorch/vision/issues/1938
+        # But, we shouldn't have to do this: https://github.com/replicate/keepsake/issues/551
+        opener = urllib.request.build_opener()
+        opener.addheaders = [("User-agent", "Mozilla/5.0")]
+        urllib.request.install_opener(opener)
+
         # download only
         MNIST(
             "/tmp/keepsake-test-mnist",