Fix has_split method

Coloquinte · Sep 17, 2022 · 44ddac1 · 44ddac1
1 parent 9b5d063
commit 44ddac1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/setup.py b/setup.py
@@ -7,7 +7,7 @@
 
 setup(
     name="torchsr",
-    version="1.0.3",
+    version="1.0.4",
     description="Super Resolution Networks for pytorch",
     long_description=long_description,
     long_description_content_type="text/markdown",

diff --git a/torchsr/datasets/common.py b/torchsr/datasets/common.py
@@ -110,15 +110,15 @@ def get_splits(cls):
         return set(sp for (t, sp, sc) in cls.track_dirs.keys())
 
     @classmethod
-    def has_split(self, split):
+    def has_split(cls, split):
         return split in cls.get_splits()
 
     def get_dir(self, track, split, scale):
         if (track, split, scale) not in self.track_dirs:
             if track not in self.get_tracks():
                 raise ValueError(f"{self.__class__.__name__} does not include track {track}. "
                                  f"Use one of {list(self.get_tracks())}")
-            if split not in self.get_splits():
+            if not self.has_split(split):
                 raise ValueError(f"{self.__class__.__name__} does not include split {split}. "
                                  f"Use one of {list(self.get_splits())}")
             available = ", ".join([str(sc) for t, sp, sc in self.track_dirs if t == track and sp == split])