fix rwkv libai code step2

libai/engine/trainer.py

@@ -279,7 +279,6 @@ def run_step(self, get_batch: Callable, input_placement_device: str = "cuda"):
         data_time = time.perf_counter() - start
 
         loss_dict = self.model(**data)
-
         losses = sum(loss_dict.values()) / self.grad_acc_steps
 
         losses.backward()
@@ -337,6 +336,7 @@ def run_step(self, get_batch: Callable, input_placement_device: str = "cuda"):
         data = get_batch(
             data, input_placement_device, getattr(self.data_loader, "mixup_func", None)
         )
+
         data_time = time.perf_counter() - start
 
         # If you want to do something with the losses, you can wrap the model.

libai/models/utils/graph_base.py

@@ -59,6 +59,8 @@ def __init__(
                     self.set_grad_scaler(grad_scaler)
                 elif amp_type == "bf16":
                     self.config.enable_amp(True, dtype=flow.bfloat16)
+                else:
+                    raise NotImplementedError
 
             if grad_acc_steps > 1:
                 self.config.set_gradient_accumulation_steps(grad_acc_steps)