Update app.py
Browse files
app.py
CHANGED
@@ -56,11 +56,11 @@ with init_empty_weights():
|
|
56 |
#
|
57 |
|
58 |
|
59 |
-
|
60 |
|
61 |
# Load the model with the inferred device map
|
62 |
# model = load_checkpoint_and_dispatch(model, model_id, device_map=device_map, no_split_module_classes=["GPTJBlock"])
|
63 |
-
model.half()
|
64 |
|
65 |
|
66 |
|
|
|
56 |
#
|
57 |
|
58 |
|
59 |
+
device_map = infer_auto_device_map(model, max_memory={0: "80GB", 1: "80GB", "cpu": "65GB"})
|
60 |
|
61 |
# Load the model with the inferred device map
|
62 |
# model = load_checkpoint_and_dispatch(model, model_id, device_map=device_map, no_split_module_classes=["GPTJBlock"])
|
63 |
+
# model.half()
|
64 |
|
65 |
|
66 |
|