Update tool/comget/generator.py
Browse files- tool/comget/generator.py +1 -1
tool/comget/generator.py
CHANGED
|
@@ -167,7 +167,7 @@ def generation(value):
|
|
| 167 |
|
| 168 |
args.model_weight = f'{args.csv_name}.pt'
|
| 169 |
model.load_state_dict(torch.load('tool/comget/' + args.model_weight))
|
| 170 |
-
model.to('
|
| 171 |
print('Model loaded')
|
| 172 |
|
| 173 |
gen_iter = math.ceil(args.gen_size / args.batch_size)
|
|
|
|
| 167 |
|
| 168 |
args.model_weight = f'{args.csv_name}.pt'
|
| 169 |
model.load_state_dict(torch.load('tool/comget/' + args.model_weight))
|
| 170 |
+
model.to('cpu')
|
| 171 |
print('Model loaded')
|
| 172 |
|
| 173 |
gen_iter = math.ceil(args.gen_size / args.batch_size)
|