badaoui HF Staff commited on
Commit
825604b
·
verified ·
1 Parent(s): dc818fc

Update optimum_neuron_export.py

Browse files
Files changed (1) hide show
  1. 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 | None = None,
121
  auto_cast_type: str = "bf16",
122
  optlevel: str = "2",
123
- instance_type: str = "trn1",
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}")