support report custom args

Former-commit-id: d41254c40a1c5cacf9377096adb27efa9bdb79ea
2024-12-19 14:57:09 +00:00
parent adff887659
commit a897d46049
20 changed files with 164 additions and 124 deletions
--- a/src/llamafactory/chat/hf_engine.py
+++ b/src/llamafactory/chat/hf_engine.py
@@ -171,7 +171,10 @@ class HuggingfaceEngine(BaseEngine):
            elif not isinstance(value, torch.Tensor):
                value = torch.tensor(value)

-            gen_kwargs[key] = value.to(dtype=model.dtype, device=model.device)
+            if torch.is_floating_point(value):
+                value = value.to(model.dtype)
+
+            gen_kwargs[key] = value.to(model.device)

        return gen_kwargs, prompt_length