fix small bug

lucidrains · Aug 31, 2021 · 461b57d · 461b57d
1 parent bead369
commit 461b57d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'tab-transformer-pytorch',
   packages = find_packages(),
-  version = '0.1.2',
+  version = '0.1.4',
   license='MIT',
   description = 'Tab Transformer - Pytorch',
   author = 'Phil Wang',

diff --git a/tab_transformer_pytorch/tab_transformer_pytorch.py b/tab_transformer_pytorch/tab_transformer_pytorch.py
@@ -112,7 +112,7 @@ def __init__(self, dims, act = None):
         dims_pairs = list(zip(dims[:-1], dims[1:]))
         layers = []
         for ind, (dim_in, dim_out) in enumerate(dims_pairs):
-            is_last = ind >= (len(dims) - 1)
+            is_last = ind >= (len(dims_pairs) - 1)
             linear = nn.Linear(dim_in, dim_out)
             layers.append(linear)