fix package/module import

freeforestml/model.py

@@ -5,10 +5,12 @@ import sys
 import h5py
 import json
 import socket
+import multiprocessing as mp
 
 import numpy as np
 import pandas as pd
 import tensorflow as tf
+import keras_tuner
 
 from freeforestml.variable import Variable
 from freeforestml.helpers import python_to_str, str_to_python
@@ -823,7 +825,7 @@ class HepNet:
         self.history.to_hdf(path, "history")
 
     @classmethod
-    def load(cls, path):
+    def load(cls, path, **kwds):
         """
         Restore a model from a hdf5 file.
         """
@@ -869,7 +871,7 @@ class HepNet:
         with h5py.File(path, "r") as input_file:
             for fold_i in range(cv.k):
                 path_token = instance._get_model_path(path, fold_i)
-                model = tf.keras.models.load_model(path_token)
+                model = tf.keras.models.load_model(path_token, **kwds)
                 instance.models.append(model)
 
         # load normalizer
@@ -1303,9 +1305,9 @@ class tuner_mp(object):
             print("Fold:%s\t ID:%s IP:%s PORT:%s"%(fold_i,os.getenv("KERASTUNER_TUNER_ID"),os.getenv("KERASTUNER_ORACLE_IP"),os.getenv("KERASTUNER_ORACLE_PORT")))
 
         #Constrain memory growth on physical GPUs
-        physical_devices = tensorflow.config.list_physical_devices('GPU')
+        physical_devices = tf.config.list_physical_devices('GPU')
         try:
-            tensorflow.config.experimental.set_memory_growth(physical_devices[0], True)
+            tf.config.experimental.set_memory_growth(physical_devices[0], True)
         except:
             # In case of CPU or virtual devices
             pass