aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorAnthony Wang2022-07-17 10:48:11 -0500
committerAnthony Wang2022-07-17 10:48:11 -0500
commit7720c1c1618d362f259b3b14b0e75946373b79b3 (patch)
tree9ac7e265084f0ed7f542aecd5ac33f1b32e28836
parente61a793dd69a1759e58e4e87be27f516c47102c6 (diff)
Update bot.py to use gpt-large
-rw-r--r--bot.py4
1 files changed, 2 insertions, 2 deletions
diff --git a/bot.py b/bot.py
index d285471..34eb039 100644
--- a/bot.py
+++ b/bot.py
@@ -18,8 +18,8 @@ parser.add_argument('-m', '--model', default='model',
args = parser.parse_args()
-tokenizer = AutoTokenizer.from_pretrained('gpt2-medium')
-model = AutoModelForCausalLM.from_pretrained(args.model, low_cpu_mem_usage=True).to('cuda')
+tokenizer = AutoTokenizer.from_pretrained('gpt2-large')
+model = AutoModelForCausalLM.from_pretrained(args.model).to('cuda')
if args.input is None: