fix: match dtype of activations during Concept.forward

chanind · Nov 23, 2023 · 3a06459 · 3a06459
1 parent 71ab816
commit 3a06459
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/README.md b/README.md
@@ -188,8 +188,8 @@ matcher = ConceptMatcher(model, tokenizer, concepts=concepts)
 
 match_info = matcher.query("Beijing is a northern city", subject="Beijing")
 
-print(match_info.best_match.name) # located in country: China
-print(match_info.betch_match.score) # 0.832
+print(match_info.best_match.concept) # located in country: China
+print(match_info.best_match.score) # 0.832
 ```
 
 ## Acknowledgements

diff --git a/docs/basic_usage.rst b/docs/basic_usage.rst
@@ -172,5 +172,5 @@ We can use the ``ConceptMatcher`` class to do this matching.
 
     match_info = matcher.query("Beijing is a northern city", subject="Beijing")
 
-    print(match_info.best_match.name) # located in country: China
-    print(match_info.betch_match.score) # 0.832
+    print(match_info.best_match.concept) # located in country: China
+    print(match_info.best_match.score) # 0.832
diff --git a/linear_relational/Concept.py b/linear_relational/Concept.py
@@ -34,7 +34,7 @@ def __init__(
         self.name = name or f"{self.relation}: {self.object}"
 
     def forward(self, activations: torch.Tensor) -> torch.Tensor:
-        vector = self.vector.to(activations.device)
+        vector = self.vector.to(activations.device, dtype=activations.dtype)
         if len(activations.shape) == 1:
             return vector @ activations
         return vector @ activations.T