diff --git a/phygnn/__init__.py b/phygnn/__init__.py
index 90829a7..9e1ded1 100644
--- a/phygnn/__init__.py
+++ b/phygnn/__init__.py
@@ -5,8 +5,12 @@
 from .base import CustomNetwork, GradientUtils
 from .phygnn import PhysicsGuidedNeuralNetwork
 from .layers import Layers, HiddenLayers
+from .layers.custom_layers import GaussianKernelInit2D
 from .utilities import PreProcess, tf_isin, tf_log10
 from phygnn.version import __version__
+from tensorflow.keras.utils import get_custom_objects
+
+get_custom_objects()['GaussianKernelInit2D'] = GaussianKernelInit2D
 
 __author__ = """Grant Buster"""
 __email__ = "grant.buster@nrel.gov"
diff --git a/tests/test_layers.py b/tests/test_layers.py
index 0a90d9b..4af14d1 100644
--- a/tests/test_layers.py
+++ b/tests/test_layers.py
@@ -477,7 +477,8 @@ def test_gaussian_kernel():
     assert kernels[1].min() > kernels[0].min()
 
     layers = [{'class': 'Conv2D', 'filters': 16, 'kernel_size': 3,
-               'kernel_initializer': GaussianKernelInit2D()}]
+               'kernel_initializer': GaussianKernelInit2D(),
+               'trainable': False}]
     model1 = TfModel.build(['a'], ['b'], hidden_layers=layers,
                            input_layer=False, output_layer=False)
     x_in = np.random.uniform(0, 1, (10, 12, 12, 1))
@@ -493,3 +494,7 @@ def test_gaussian_kernel():
         out2 = model2.predict(x_in)
         assert np.allclose(kernel1, kernel2)
         assert np.allclose(out1, out2)
+
+        layer = model2.layers[0]
+        x_in = np.random.uniform(0, 1, (10, 24, 24, 1))
+        _ = model2.predict(x_in)