add _autocast_context_manager to visu.py and solver.predict

Author: HydrogenSulfate

ppsci/solver/solver.py

@@ -472,10 +472,7 @@ def predict(
                 batch_input_dict[key].stop_gradient = False
 
             # forward
-            if self.use_amp:
-                with amp.auto_cast(level=self.amp_level):
-                    batch_output_dict = self.model(batch_input_dict)
-            else:
+            with self._autocast_context_manager():
                 batch_output_dict = self.model(batch_input_dict)
 
             # collect batch data

ppsci/solver/visu.py

@@ -65,10 +65,7 @@ def visualize_func(solver, epoch_id: int):
                 evaluator.add_target_expr(output_expr, output_key)
 
             # forward
-            if solver.use_amp:
-                with amp.auto_cast(level=solver.amp_level):
-                    batch_output_dict = evaluator(batch_input_dict)
-            else:
+            with solver._autocast_context_manager():
                 batch_output_dict = evaluator(batch_input_dict)
 
             # collect batch data