Update GeMOSA v4

2021-05-27 17:46:00 +08:00 · 2021-05-27 17:46:00 +08:00 · 16861f0f3d
commit 16861f0f3d
parent b6e11c6360
1 changed files with 4 additions and 0 deletions
--- a/xautodl/datasets/synthetic_env.py
+++ b/xautodl/datasets/synthetic_env.py
@ -127,6 +127,10 @@ class SyntheticDEnv(data.Dataset):
            targets = torch.from_numpy(targets)
        else:
            targets = torch.Tensor(targets)
+        if dataset.dtype == torch.float64:
+            dataset = dataset.float()
+        if targets.dtype == torch.float64:
+            targets = targets.float()
        return torch.Tensor([timestamp]), (dataset, targets)

    def __len__(self):