diff --git a/python/pytest/conftest.py b/python/pytest/conftest.py
index 766a3cb..b25600a 100644
--- a/python/pytest/conftest.py
+++ b/python/pytest/conftest.py
@@ -2,9 +2,11 @@
 
 import pytest
 
+
 @pytest.fixture(scope="session")
 def ptens_cuda_support():
     import ptens_base
+
     string = ptens_base.status_str().split("\n")
     for line in string:
         if "CUDA support" in line:
@@ -14,18 +16,20 @@ def ptens_cuda_support():
                 return False
     assert False
 
+
 @pytest.fixture(scope="session")
 def device(ptens_cuda_support):
-    
     device = os.environ["TORCH_TEST_DEVICE"]
 
     if "cuda" in device:
         assert ptens_cuda_support
         import torch
+
         assert torch.cuda.is_available()
 
     return device
 
+
 @pytest.fixture(scope="session")
 def float_epsilon():
     return 1e-5
diff --git a/python/pytest/test_doc.py b/python/pytest/test_doc.py
index 2137f6f..c255b3f 100644
--- a/python/pytest/test_doc.py
+++ b/python/pytest/test_doc.py
@@ -1,18 +1,19 @@
 import torch
 import ptens
 
+
 def test_create_tensor(device):
-    A = torch.ones((3,5))
+    A = torch.ones((3, 5))
     sum_a = float(torch.sum(A**2))
-    assert A.shape == (3,5)
+    assert A.shape == (3, 5)
     B = A.to(device)
     sum_b = float(torch.sum(B**2))
     assert device in str(B.device)
     assert abs(sum_a - sum_b) < 1e-6
-        
+
 
 def test_create_ptensor0(device, float_epsilon):
-    A = ptens.ptensor0.randn([2],5)
+    A = ptens.ptensor0.randn([2], 5)
     sum_a = float(torch.sum(A**2))
     B = A.to(device)
     sum_b = float(torch.sum(B**2))
@@ -23,37 +24,38 @@ def test_create_ptensor0(device, float_epsilon):
     assert sum_a > 1e-3
     assert abs(sum_a - sum_b) < float_epsilon
 
+
 def test_create_ptensor1(device, float_epsilon):
-    A = ptens.ptensor1.randn([1,2,3],5)
+    A = ptens.ptensor1.randn([1, 2, 3], 5)
     sum_a = float(torch.sum(A**2))
     B = A.to(device)
     sum_b = float(torch.sum(B**2))
     print(B)
-    assert B.atoms == [1,2,3]
-    assert B.shape == (3,5)
+    assert B.atoms == [1, 2, 3]
+    assert B.shape == (3, 5)
     assert device in str(B.device)
     assert abs(sum_a - sum_b) < float_epsilon
     assert sum_a > 1e-3
 
 
 def test_create_ptensor2(device, float_epsilon):
-    A=ptens.ptensor2.randn([1,2,3],5)
+    A = ptens.ptensor2.randn([1, 2, 3], 5)
     sum_a = float(torch.sum(A**2))
     B = A.to(device)
     sum_b = float(torch.sum(B**2))
     print(B)
-    assert B.atoms == [1,2,3]
-    assert B.shape == (3,3,5)
+    assert B.atoms == [1, 2, 3]
+    assert B.shape == (3, 3, 5)
     assert device in str(B.device)
     assert abs(sum_a - sum_b) < float_epsilon
     assert sum_a > 1e-3
 
+
 def test_sequential(device, float_epsilon):
-    A = ptens.ptensor1.sequential([1,2,3],5)
+    A = ptens.ptensor1.sequential([1, 2, 3], 5)
     sum_a = float(torch.sum(A))
     B = A.to(device)
     sum_b = float(torch.sum(B))
 
     assert abs(sum_a - sum_b) < float_epsilon
-    assert abs(sum_a - sum(range(3*5))) < float_epsilon
-    
+    assert abs(sum_a - sum(range(3 * 5))) < float_epsilon
diff --git a/python/pytest/test_tests.py b/python/pytest/test_tests.py
index 0060b2c..9c1955f 100644
--- a/python/pytest/test_tests.py
+++ b/python/pytest/test_tests.py
@@ -3,16 +3,16 @@
 from ptens_base import atomspack
 from torch.autograd.gradcheck import gradcheck, gradgradcheck
 
+
 def test_bug1(device):
     nnodes = 15
-    graph =ptens.ggraph.random(nnodes,0.5)
+    graph = ptens.ggraph.random(nnodes, 0.5)
     print(graph)
     subgraphs = [ptens.subgraph.trivial(), ptens.subgraph.edge()]
     node_values = torch.rand(nnodes, 1, requires_grad=True)
 
     node_attributes = ptens.subgraphlayer0.from_matrix(graph, ptens.subgraph.trivial(), node_values)
 
-
     for sg in subgraphs:
         gather_features = ptens.subgraphlayer0.gather(sg, node_attributes)
         result = torch.sum(gather_features)
@@ -26,4 +26,3 @@ def test_bug1(device):
 
         check = gradcheck(ptens.subgraphlayer0.gather, (sg, node_attributes), eps=1e-3)
         print(check)
-
diff --git a/python/pytest/test_tests4.py b/python/pytest/test_tests4.py
index 4629ea8..b8a175d 100644
--- a/python/pytest/test_tests4.py
+++ b/python/pytest/test_tests4.py
@@ -2,33 +2,32 @@
 import ptens_base as pb
 import ptens as p
 
+
 def test_ptensorlayer(device, float_epsilon):
+    atoms = pb.atomspack.from_list([[1, 3, 4], [2, 5], [0, 2]])
+    atoms2 = pb.atomspack.random(5, 5, 0.6)
 
-    atoms=pb.atomspack.from_list([[1,3,4],[2,5],[0,2]])
-    atoms2=pb.atomspack.random(5,5,0.6)
+    A0 = p.ptensorlayer0.randn(atoms, 3)
+    A1 = p.ptensorlayer1.randn(atoms, 3)
+    A2 = p.ptensorlayer2.randn(atoms, 3)
 
-    A0=p.ptensorlayer0.randn(atoms,3)
-    A1=p.ptensorlayer1.randn(atoms,3)
-    A2=p.ptensorlayer2.randn(atoms,3)
+    B0 = p.ptensorlayer1.gather(atoms2, A0)
+    B1 = p.ptensorlayer1.gather(atoms2, A1)
+    B2 = p.ptensorlayer1.gather(atoms2, A2)
 
-    B0=p.ptensorlayer1.gather(atoms2,A0)
-    B1=p.ptensorlayer1.gather(atoms2,A1)
-    B2=p.ptensorlayer1.gather(atoms2,A2)
-    
-    A0g=A0.to(device)
-    A1g=A1.to(device)
-    A2g=A2.to(device)
+    A0g = A0.to(device)
+    A1g = A1.to(device)
+    A2g = A2.to(device)
 
-    B0g=p.ptensorlayer1.gather(atoms2,A0g)
-    B1g=p.ptensorlayer1.gather(atoms2,A1g)
-    B2g=p.ptensorlayer1.gather(atoms2,A2g)
+    B0g = p.ptensorlayer1.gather(atoms2, A0g)
+    B1g = p.ptensorlayer1.gather(atoms2, A1g)
+    B2g = p.ptensorlayer1.gather(atoms2, A2g)
 
     # Making sure we don't just compute on zeros
     assert torch.norm(B0) > float_epsilon
     assert torch.norm(B1) > float_epsilon
     assert torch.norm(B2) > float_epsilon
-    
-    assert torch.allclose(B0, B0g.to("cpu"))
-    assert torch.allclose(B1g.to('cpu'), B1)
-    assert torch.allclose(B2g.to('cpu'), B2)
 
+    assert torch.allclose(B0, B0g.to("cpu"))
+    assert torch.allclose(B1g.to("cpu"), B1)
+    assert torch.allclose(B2g.to("cpu"), B2)