From 7e4fba9d10c299ddd78427be3d9736b4bf9bce96 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Sun, 26 Apr 2020 22:54:49 +0300
Subject: [PATCH] Fix scikit-learn module detection

---
 trains/backend_interface/task/repo/scriptinfo.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/trains/backend_interface/task/repo/scriptinfo.py b/trains/backend_interface/task/repo/scriptinfo.py
index ee1c23eb..ab71f459 100644
--- a/trains/backend_interface/task/repo/scriptinfo.py
+++ b/trains/backend_interface/task/repo/scriptinfo.py
@@ -60,6 +60,13 @@ class ScriptRequirements(object):
             modules.add('azure_storage_blob', 'trains.storage', 0)
         except Exception:
             pass
+
+        # bugfix, replace sklearn with scikit-learn name
+        if 'sklearn' in modules:
+            sklearn = modules.pop('sklearn', {})
+            for fname, lines in sklearn.items():
+                modules.add('scikit_learn', fname, lines)
+
         # if we have torch and it supports tensorboard, we should add that as well
         # (because it will not be detected automatically)
         if 'torch' in modules and 'tensorboard' not in modules: