tiny fix

Former-commit-id: d2cede7023bbe28525ef8b4ad27247445d8c22e5

src/llamafactory/extras/misc.py

@@ -156,6 +156,18 @@ def get_logits_processor() -> "LogitsProcessorList":
     return logits_processor
 
 
+def get_peak_memory() -> Tuple[int, int]:
+    r"""
+    Gets the peak memory usage for the current device (in Bytes).
+    """
+    if is_torch_npu_available():
+        return torch.npu.max_memory_allocated(), torch.npu.max_memory_reserved()
+    elif is_torch_cuda_available():
+        return torch.cuda.max_memory_allocated(), torch.cuda.max_memory_reserved()
+    else:
+        return 0, 0
+
+
 def has_tokenized_data(path: "os.PathLike") -> bool:
     r"""
     Checks if the path has a tokenized dataset.