diff --git a/trains/binding/frameworks/lightgbm_bind.py b/trains/binding/frameworks/lightgbm_bind.py
index b2f7c94d..e8bf116c 100644
--- a/trains/binding/frameworks/lightgbm_bind.py
+++ b/trains/binding/frameworks/lightgbm_bind.py
@@ -118,11 +118,12 @@ class PatchLIGHTgbmModelIO(PatchBaseModelIO):
                 except Exception:
                     pass
             return callback
-        params, train_set = args
+        
         kwargs.setdefault("callbacks", []).append(trains_lightgbm_callback())
-        ret = original_fn(params, train_set, **kwargs)
+        ret = original_fn(*args, **kwargs)
         if not PatchLIGHTgbmModelIO.__main_task:
             return ret
+        params = args[0] if args else kwargs['params']
         for k, v in params.items():
             if isinstance(v, set):
                 params[k] = list(v)