Michael Gira commited on
Commit
659007c
1 Parent(s): 67fdebe

Change device

Browse files
Files changed (1) hide show
  1. load_model.py +2 -1
load_model.py CHANGED
@@ -4,11 +4,12 @@ import torch
4
  from transformers import GPT2Tokenizer, GPT2LMHeadModel
5
  from model import get_model
6
 
7
- device = 'cuda'
8
  models_path = 'models'
9
 
10
  tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
11
 
 
12
  def load_gpt2():
13
 
14
  model = GPT2LMHeadModel.from_pretrained('gpt2').to(device)
4
  from transformers import GPT2Tokenizer, GPT2LMHeadModel
5
  from model import get_model
6
 
7
+ device = 'cpu'
8
  models_path = 'models'
9
 
10
  tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
11
 
12
+
13
  def load_gpt2():
14
 
15
  model = GPT2LMHeadModel.from_pretrained('gpt2').to(device)