replace cuda for jit scripted modules

Signed-off-by: Nithin Rao Koluguri <nithinraok>

nemo/utils/export_utils.py

@@ -149,7 +149,7 @@ def verify_torchscript(model, output, input_examples, check_tolerance=0.01):
     for input_example in input_examples:
         input_list, input_dict = parse_input_example(input_example)
         # We disable autocast here to make sure exported TS will run under Triton or other C++ env
-        with torch.amp.autocast(model.device.type, enabled=False):
+        with torch.amp.autocast('cuda', enabled=False):
             output_example = model.forward(*input_list, **input_dict)
             ts_model = torch.jit.load(output)
             all_good = all_good and run_ts_and_compare(