File tree 4 files changed +4
-0
lines changed
4 files changed +4
-0
lines changed Original file line number Diff line number Diff line change @@ -172,6 +172,7 @@ def main():
172
172
jit_config = None
173
173
if args .use_torch_jit :
174
174
from nemo .automodel .compiler import TorchCompileConfig
175
+
175
176
jit_config = TorchCompileConfig (kwargs = {'dynamic' : True })
176
177
177
178
if args .use_te_optimizer :
Original file line number Diff line number Diff line change @@ -171,6 +171,7 @@ def main():
171
171
jit_config = None
172
172
if args .use_torch_jit :
173
173
from nemo .automodel .compiler import TorchCompileConfig
174
+
174
175
jit_config = TorchCompileConfig (kwargs = {'dynamic' : True })
175
176
176
177
if args .use_te_optimizer :
Original file line number Diff line number Diff line change @@ -281,6 +281,7 @@ def main():
281
281
jit_config = None
282
282
if args .use_torch_jit :
283
283
from nemo .automodel .compiler import TorchCompileConfig
284
+
284
285
jit_config = TorchCompileConfig (kwargs = {'dynamic' : True })
285
286
286
287
model = llm .HFAutoModelForCausalLM (model_name = args .model , compiler_config = jit_config )
Original file line number Diff line number Diff line change @@ -324,6 +324,7 @@ def main():
324
324
jit_config = None
325
325
if args .use_torch_jit :
326
326
from nemo .automodel .compiler import TorchCompileConfig
327
+
327
328
jit_config = TorchCompileConfig (kwargs = {'dynamic' : True })
328
329
329
330
if args .auto_resume :
You can’t perform that action at this time.
0 commit comments