experiment training

JVGD · Nov 1, 2021 · eebaf4e · eebaf4e
1 parent 8a35608
commit eebaf4e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/experiment1/dataset.py b/experiment1/dataset.py
@@ -7,7 +7,7 @@ class NumberAdd(Dataset):
     def __init__(self, total_samples: int) -> None:
         """Generate dataset with pairs of numbers and their sum"""
         self.N = total_samples
-        self.samples = T.randint(low=-1000, high=1000, size=(self.N, 2))
+        self.samples = T.randint(low=-1000, high=1000, size=(self.N, 2), dtype=T.float32)
         self.targets = T.sum(self.samples, dim=-1)
 
     def __getitem__(self, index: int) -> tuple:

diff --git a/experiment1/modules.py b/experiment1/modules.py
@@ -8,7 +8,7 @@ def __init__(self) -> None:
 
     def __call__(self, X: T.Tensor) -> T.Tensor:
         """Calling the adding operation without gradient"""
-        return X[0] + X[1]
+        return X[:, 0] + X[:, 1]
 
 
 class Substracter(object):
@@ -18,7 +18,7 @@ def __init__(self) -> None:
 
     def __call__(self, X: T.Tensor) -> T.Tensor:
         """Calling the adding operation without gradient"""
-        return X[0] - X[1]
+        return X[:, 0] - X[:, 1]
 
 
 if __name__ == "__main__":