Merge pull request #109 from BlackSamorez/sharded_tests

Testing interfaces (soon to be refactored)
BlackSamorez · Jul 25, 2023 · aa3dd94 · aa3dd94
2 parents 31329ad + b915ad8
commit aa3dd94
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/tests/test_legacy.py b/tests/test_legacy.py
@@ -0,0 +1,14 @@
+import torch
+from transformers import BertModel, PreTrainedModel
+
+from tensor_parallel import Sharded, tensor_parallel
+
+
+def test_legacy_factory_and_sharded():
+    model = BertModel.from_pretrained("bert-base-uncased")
+
+    tp_model = tensor_parallel(model, sharded=False)
+    assert isinstance(tp_model, PreTrainedModel)
+    tp_model.wrapped_model = Sharded(tp_model.wrapped_model)
+
+    tp_model(torch.zeros(1, 8, dtype=int))