Implementing __contains__ for Vectors class (#2144)

pytorch · Apr 6, 2023 · f3c34f9 · f3c34f9
1 parent fdc4858
commit f3c34f9
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/torchtext/vocab/vectors.py b/torchtext/vocab/vectors.py
@@ -64,6 +64,9 @@ def __getitem__(self, token):
         else:
             return self.unk_init(torch.Tensor(self.dim))
 
+    def __contains__(self, token):
+        return token in self.stoi
+
     def cache(self, name, cache, url=None, max_vectors=None):
         import ssl