add device
Browse files- app.py +2 -0
- app_test.py +2 -0
app.py
CHANGED
|
@@ -330,6 +330,8 @@ def construction_layout():
|
|
| 330 |
print("after .to(device)")
|
| 331 |
model = model.bfloat16()
|
| 332 |
model.eval()
|
|
|
|
|
|
|
| 333 |
return model, quantizer, tokenizer, params_dict["width"], params_dict["height"], device
|
| 334 |
|
| 335 |
@torch.no_grad()
|
|
|
|
| 330 |
print("after .to(device)")
|
| 331 |
model = model.bfloat16()
|
| 332 |
model.eval()
|
| 333 |
+
quantizer = quantizer.to("cuda")
|
| 334 |
+
tokenizer = tokenizer.to("cuda")
|
| 335 |
return model, quantizer, tokenizer, params_dict["width"], params_dict["height"], device
|
| 336 |
|
| 337 |
@torch.no_grad()
|
app_test.py
CHANGED
|
@@ -330,6 +330,8 @@ def construction_layout():
|
|
| 330 |
print("after .to(device)")
|
| 331 |
model = model.bfloat16()
|
| 332 |
model.eval()
|
|
|
|
|
|
|
| 333 |
return model, quantizer, tokenizer, params_dict["width"], params_dict["height"], device
|
| 334 |
|
| 335 |
@torch.no_grad()
|
|
|
|
| 330 |
print("after .to(device)")
|
| 331 |
model = model.bfloat16()
|
| 332 |
model.eval()
|
| 333 |
+
quantizer = quantizer.to("cuda")
|
| 334 |
+
tokenizer = tokenizer.to("cuda")
|
| 335 |
return model, quantizer, tokenizer, params_dict["width"], params_dict["height"], device
|
| 336 |
|
| 337 |
@torch.no_grad()
|