Update.

[pytorch.git] / gpt-test.py

diff --git a/gpt-test.py b/gpt-test.py
index 557f734..0967043 100755 (executable)
--- a/gpt-test.py
+++ b/gpt-test.py
@@ -53,6 +53,8 @@ tokenizer = GPT2Tokenizer.from_pretrained(model_name)
 model = GPT2LMHeadModel.from_pretrained(model_name)
 model.eval()
 
+print(f'Using {model_name} ({int(sum(p.numel() for p in model.parameters())/(1e6))}M parameters)')
+
 print(
     complete(model, tokenizer,
              'The object was blue all over, but also green all over, it was a',