Fix scikit-learn module detection

This commit is contained in:
allegroai 2020-04-26 22:54:49 +03:00
parent 13ce783fa3
commit 7e4fba9d10

View File

@ -60,6 +60,13 @@ class ScriptRequirements(object):
modules.add('azure_storage_blob', 'trains.storage', 0) modules.add('azure_storage_blob', 'trains.storage', 0)
except Exception: except Exception:
pass 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 # if we have torch and it supports tensorboard, we should add that as well
# (because it will not be detected automatically) # (because it will not be detected automatically)
if 'torch' in modules and 'tensorboard' not in modules: if 'torch' in modules and 'tensorboard' not in modules: