polish qtransformer

opendilab · Jun 20, 2024 · 7095b38 · 7095b38
1 parent 6023c65
commit 7095b38
Show file tree

Hide file tree

Showing 4 changed files with 1,155 additions and 634 deletions.
diff --git a/dataset/qtransformer.py b/dataset/qtransformer.py
@@ -8,6 +8,7 @@
 from rich.progress import track
 from torch.utils.data import DataLoader, Dataset
 
+
 # just force training on 64 bit systems
 
 assert sys.maxsize > (
@@ -33,7 +34,6 @@ def cast_tuple(t):
 
 # replay memory dataset
 class ReplayMemoryDataset(Dataset):
-    @beartype
     def __init__(self, config):
         dataset_folder = config.dataset_folder
         num_timesteps = config.num_timesteps