Spaces:
Running
Running
Update optimum_neuron_export.py
Browse files- optimum_neuron_export.py +2 -2
optimum_neuron_export.py
CHANGED
|
@@ -117,10 +117,10 @@ def get_default_inputs(task_or_pipeline: str) -> Dict[str, int]:
|
|
| 117 |
return {"batch_size": 1, "sequence_length": 128}
|
| 118 |
|
| 119 |
def prepare_compiler_flags(
|
| 120 |
-
auto_cast: str
|
| 121 |
auto_cast_type: str = "bf16",
|
| 122 |
optlevel: str = "2",
|
| 123 |
-
instance_type: str = "
|
| 124 |
):
|
| 125 |
if auto_cast is not None:
|
| 126 |
logger.info(f"Using Neuron: --auto-cast {auto_cast}")
|
|
|
|
| 117 |
return {"batch_size": 1, "sequence_length": 128}
|
| 118 |
|
| 119 |
def prepare_compiler_flags(
|
| 120 |
+
auto_cast: str = "matmul",
|
| 121 |
auto_cast_type: str = "bf16",
|
| 122 |
optlevel: str = "2",
|
| 123 |
+
instance_type: str = "inf2",
|
| 124 |
):
|
| 125 |
if auto_cast is not None:
|
| 126 |
logger.info(f"Using Neuron: --auto-cast {auto_cast}")
|