Minor fix in checkpoint saving

.gitignore

@@ -0,0 +1,2 @@
+*.egg-info/
+__pycache__/

train.py → deap/train.py

@@ -35,7 +35,7 @@ def train(cfg, device='cuda'):
             model.train()
 
             opt.zero_grad()
-            
+
             outputs = model(sample)
             loss, _ = task.loss_function(outputs, sample)
 
@@ -46,10 +46,16 @@ def train(cfg, device='cuda'):
                 continue_training = False
                 break
 
-            iteration += 1   
-
-    torch.save({k:p for k,p in model.state_dict().items() if p.requires_grad}, cfg['name'] + '-weights.pth')
+            iteration += 1
 
+    torch.save(
+        {
+            name: p.detach().cpu()
+            for name, p in model.named_parameters()
+            if p.requires_grad
+        },
+        cfg['name'] + '-weights.pth'
+    )
 
 def evaluate(cfg, weights, device='cuda'):
 

pyproject.toml

@@ -0,0 +1,28 @@
+[build-system]
+requires = ["setuptools>=42.0.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+exclude = [
+    ".bzr",
+    ".direnv",
+    ".eggs",
+    ".git",
+    ".mypy_cache",
+    ".pants.d",
+    ".ruff_cache",
+    ".svn",
+    ".tox",
+    ".venv",
+    "__pypackages__",
+    "_build",
+    "buck-out",
+    "build",
+    "dist",
+    "node_modules",
+    "venv",
+    "*vendored*",
+    "*_vendor*",
+]
+
+target-version = "py38"
+fix = true

setup.cfg

@@ -0,0 +1,17 @@
+[metadata]
+name = deap
+version = 0.0.1
+description = Dense Attentive Probing
+url = https://github.com/timojl/deap
+author = Timo Lueddecke
+author_email = yourname@example.com
+
+[options]
+packages = find:
+python_requires = >=3.10
+include_package_data = True
+package_dir =
+    = .
+
+[options.packages.find]
+where = .