Fix dtype conversion bug

Bitbol-Lab · Dec 2, 2023 · b66dd59 · b66dd59
1 parent 73e6239
commit b66dd59
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/diffpass/train.py b/diffpass/train.py
@@ -247,7 +247,7 @@ def _fit(
                     )
                     results.hard_perms.append(
                         [
-                            _dccn(perms_this_group).argmax(-1).to(torch.int16)
+                            _dccn(perms_this_group.argmax(-1).to(torch.int16))
                             for perms_this_group in perms
                         ]
                     )
@@ -653,7 +653,7 @@ def _fit(
                     )
                     results.hard_perms.append(
                         [
-                            _dccn(perms_this_group).argmax(-1).to(torch.int16)
+                            _dccn(perms_this_group.argmax(-1).to(torch.int16))
                             for perms_this_group in perms
                         ]
                     )
@@ -1024,7 +1024,7 @@ def _fit(
                     )
                     results.hard_perms.append(
                         [
-                            _dccn(perms_this_group).argmax(-1).to(torch.int16)
+                            _dccn(perms_this_group.argmax(-1).to(torch.int16))
                             for perms_this_group in perms
                         ]
                     )

diff --git a/nbs/train.ipynb b/nbs/train.ipynb
@@ -319,7 +319,7 @@
     "                    )\n",
     "                    results.hard_perms.append(\n",
     "                        [\n",
-    "                            _dccn(perms_this_group).argmax(-1).to(torch.int16)\n",
+    "                            _dccn(perms_this_group.argmax(-1).to(torch.int16))\n",
     "                            for perms_this_group in perms\n",
     "                        ]\n",
     "                    )\n",
@@ -733,7 +733,7 @@
     "                    )\n",
     "                    results.hard_perms.append(\n",
     "                        [\n",
-    "                            _dccn(perms_this_group).argmax(-1).to(torch.int16)\n",
+    "                            _dccn(perms_this_group.argmax(-1).to(torch.int16))\n",
     "                            for perms_this_group in perms\n",
     "                        ]\n",
     "                    )\n",
@@ -1112,7 +1112,7 @@
     "                    )\n",
     "                    results.hard_perms.append(\n",
     "                        [\n",
-    "                            _dccn(perms_this_group).argmax(-1).to(torch.int16)\n",
+    "                            _dccn(perms_this_group.argmax(-1).to(torch.int16))\n",
     "                            for perms_this_group in perms\n",
     "                        ]\n",
     "                    )\n",