diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml
new file mode 100644
index 0000000..18bd3dc
--- /dev/null
+++ b/.github/workflows/run_tests.yml
@@ -0,0 +1,69 @@
+name: Run Tests
+
+on:
+  push:
+    tags:
+      - '*'
+    branches:
+      - '**'
+  pull_request:
+    types: [opened, synchronize, reopened, ready_for_review]
+
+jobs:
+  cancel_previous_runs:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Cancel Previous Runs
+        uses: styfle/cancel-workflow-action@0.11.0
+        with:
+          access_token: ${{ github.token }}
+
+  run_tests:
+    needs: cancel_previous_runs
+    runs-on: ubuntu-latest
+
+    strategy:
+      fail-fast: false
+      matrix:
+        config:
+        - {python: "3.7", tensorflow: "2.0"}
+        - {python: "3.7", tensorflow: "2.1"}
+        - {python: "3.8", tensorflow: "2.2"}
+        - {python: "3.8", tensorflow: "2.3"}
+        - {python: "3.8", tensorflow: "2.4"}
+        - {python: "3.9", tensorflow: "2.5"}
+        - {python: "3.9", tensorflow: "2.7"}
+        - {python: "3.9", tensorflow: "2.8"}
+        - {python: "3.9", tensorflow: "2.9"}
+        - {python: "3.10", tensorflow: "2.10"}
+
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+
+    - name: Set up environment
+      run: |
+        echo "LINUX_VERSION=$(uname -rs)" >> $GITHUB_ENV
+
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.config.python }}
+
+    - name: Cache Python packages
+      uses: actions/cache@v3
+      with:
+        path: ${{ env.pythonLocation }}
+        key: ${{env.LINUX_VERSION}}-pip-${{ matrix.config.python }}-${{ hashFiles('setup.py') }}
+        restore-keys: ${{env.LINUX_VERSION}}-pip-${{ matrix.config.python }}-
+
+    - name: Install package & dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -U wheel setuptools
+        pip install -U .[test] tensorflow==${{ matrix.config.tensorflow }}
+        python -c "import dca"
+
+    - name: Run tests
+      run: pytest -vv
diff --git a/.gitignore b/.gitignore
index a303696..2627f41 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,4 +5,6 @@ build
 *.egg-info
 .Rproj.user
 docs/build
-data/simulation/
+data/
+.idea
+
diff --git a/dca/__init__.py b/dca/__init__.py
index f815d69..e69de29 100644
--- a/dca/__init__.py
+++ b/dca/__init__.py
@@ -1,2 +0,0 @@
-import os
-os.environ['KERAS_BACKEND'] = 'tensorflow'
diff --git a/dca/__main__.py b/dca/__main__.py
index 0acbff5..3a7707c 100644
--- a/dca/__main__.py
+++ b/dca/__main__.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-import os, sys, argparse
+import argparse
 
 def parse_args():
     parser = argparse.ArgumentParser(description='Autoencoder')
diff --git a/dca/api.py b/dca/api.py
index 9401b78..5fc7aa8 100644
--- a/dca/api.py
+++ b/dca/api.py
@@ -1,4 +1,4 @@
-import os, tempfile, shutil, random
+import os, random
 import anndata
 import numpy as np
 import scanpy as sc
diff --git a/dca/hyper.py b/dca/hyper.py
index 72e6d0d..8518680 100644
--- a/dca/hyper.py
+++ b/dca/hyper.py
@@ -5,7 +5,7 @@
 import numpy as np
 from kopt import CompileFN, test_fn
 from hyperopt import fmin, tpe, hp, Trials
-import keras.optimizers as opt
+import tensorflow.keras.optimizers as opt
 
 from . import io
 from .network import AE_types
diff --git a/dca/io.py b/dca/io.py
index 1e3f47b..8b15f92 100644
--- a/dca/io.py
+++ b/dca/io.py
@@ -18,14 +18,13 @@
 from __future__ import division
 from __future__ import print_function
 
-import pickle, os, numbers
+import pickle
 
 import numpy as np
 import scipy as sp
 import pandas as pd
 import scanpy as sc
 from sklearn.model_selection import train_test_split
-from sklearn.preprocessing import scale
 
 
 #TODO: Fix this
diff --git a/dca/layers.py b/dca/layers.py
index c5ed7be..9e4ce1b 100644
--- a/dca/layers.py
+++ b/dca/layers.py
@@ -1,7 +1,6 @@
-from keras.engine.topology import Layer
-from keras.layers import Lambda, Dense
-from keras.engine.base_layer import InputSpec
-from keras import backend as K
+from tensorflow.keras.layers import Layer, Lambda, Dense
+from tensorflow.keras.layers import InputSpec
+from tensorflow.keras import backend as K
 import tensorflow as tf
 
 
@@ -81,5 +80,15 @@ def call(self, inputs):
         return output
 
 
-nan2zeroLayer = Lambda(lambda x: tf.where(tf.is_nan(x), tf.zeros_like(x), x))
-ColwiseMultLayer = Lambda(lambda l: l[0]*tf.reshape(l[1], (-1,1)))
+class ColwiseMultLayer(Layer):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+    def call(self, l):
+        return l[0]*tf.reshape(l[1], (-1,1))
+
+    def compute_output_shape(self, input_shape):
+        return input_shape[0]
diff --git a/dca/loss.py b/dca/loss.py
index e4ea49f..8331932 100644
--- a/dca/loss.py
+++ b/dca/loss.py
@@ -1,6 +1,5 @@
 import numpy as np
 import tensorflow as tf
-from keras import backend as K
 
 
 def _nan2zero(x):
diff --git a/dca/network.py b/dca/network.py
index 64a24ae..dc33269 100644
--- a/dca/network.py
+++ b/dca/network.py
@@ -15,18 +15,14 @@
 
 import os
 import pickle
-from abc import ABCMeta, abstractmethod
 
 import numpy as np
-import scanpy as sc
 
-import keras
-from keras.layers import Input, Dense, Dropout, Activation, BatchNormalization, Lambda
-from keras.models import Model
-from keras.regularizers import l1_l2
-from keras.objectives import mean_squared_error
-from keras.initializers import Constant
-from keras import backend as K
+from tensorflow.keras.layers import Input, Dense, Dropout, Activation, BatchNormalization, Lambda
+from tensorflow.keras.models import Model
+from tensorflow.keras.regularizers import l1_l2
+from tensorflow.keras.losses import mean_squared_error
+from tensorflow.keras import backend as K
 
 import tensorflow as tf
 
@@ -130,7 +126,7 @@ def build(self):
             # Use separate act. layers to give user the option to get pre-activations
             # of layers when requested
             if self.activation in advanced_activations:
-                last_hidden = keras.layers.__dict__[self.activation](name='%s_act'%layer_name)(last_hidden)
+                last_hidden = tf.keras.layers.__dict__[self.activation](name='%s_act'%layer_name)(last_hidden)
             else:
                 last_hidden = Activation(self.activation, name='%s_act'%layer_name)(last_hidden)
 
@@ -146,7 +142,7 @@ def build_output(self):
         mean = Dense(self.output_size, kernel_initializer=self.init,
                      kernel_regularizer=l1_l2(self.l1_coef, self.l2_coef),
                      name='mean')(self.decoder_output)
-        output = ColwiseMultLayer([mean, self.sf_layer])
+        output = ColwiseMultLayer()([mean, self.sf_layer])
 
         # keep unscaled output as an extra model
         self.extra_models['mean_norm'] = Model(inputs=self.input_layer, outputs=mean)
@@ -236,7 +232,7 @@ def build_output(self):
         mean = Dense(self.output_size, activation=MeanAct, kernel_initializer=self.init,
                      kernel_regularizer=l1_l2(self.l1_coef, self.l2_coef),
                      name='mean')(self.decoder_output)
-        output = ColwiseMultLayer([mean, self.sf_layer])
+        output = ColwiseMultLayer()([mean, self.sf_layer])
         self.loss = poisson_loss
 
         self.extra_models['mean_norm'] = Model(inputs=self.input_layer, outputs=mean)
@@ -257,7 +253,7 @@ def build_output(self):
         disp = ConstantDispersionLayer(name='dispersion')
         mean = disp(mean)
 
-        output = ColwiseMultLayer([mean, self.sf_layer])
+        output = ColwiseMultLayer()([mean, self.sf_layer])
 
         nb = NB(disp.theta_exp)
         self.loss = nb.loss
@@ -302,7 +298,7 @@ def build_output(self):
         mean = Dense(self.output_size, activation=MeanAct, kernel_initializer=self.init,
                        kernel_regularizer=l1_l2(self.l1_coef, self.l2_coef),
                        name='mean')(self.decoder_output)
-        output = ColwiseMultLayer([mean, self.sf_layer])
+        output = ColwiseMultLayer()([mean, self.sf_layer])
         output = SliceLayer(0, name='slice')([output, disp])
 
         nb = NB(theta=disp, debug=self.debug)
@@ -350,7 +346,7 @@ def build_output(self):
         mean = Dense(self.output_size, activation=MeanAct, kernel_initializer=self.init,
                        kernel_regularizer=l1_l2(self.l1_coef, self.l2_coef),
                        name='mean')(self.decoder_output)
-        output = ColwiseMultLayer([mean, self.sf_layer])
+        output = ColwiseMultLayer()([mean, self.sf_layer])
         output = SliceLayer(0, name='slice')([output, disp])
 
         nb = NB(theta=disp, debug=self.debug)
@@ -378,7 +374,7 @@ def build_output(self):
         mean = Dense(self.output_size, activation=MeanAct, kernel_initializer=self.init,
                        kernel_regularizer=l1_l2(self.l1_coef, self.l2_coef),
                        name='mean')(self.decoder_output)
-        output = ColwiseMultLayer([mean, self.sf_layer])
+        output = ColwiseMultLayer()([mean, self.sf_layer])
         output = SliceLayer(0, name='slice')([output, disp, pi])
 
         zinb = ZINB(pi, theta=disp, ridge_lambda=self.ridge, debug=self.debug)
@@ -446,7 +442,7 @@ def build_output(self):
 
         mean = Activation(MeanAct, name='mean')(mean_no_act)
 
-        output = ColwiseMultLayer([mean, self.sf_layer])
+        output = ColwiseMultLayer()([mean, self.sf_layer])
         output = SliceLayer(0, name='slice')([output, disp, pi])
 
         zinb = ZINB(pi, theta=disp, ridge_lambda=self.ridge, debug=self.debug)
@@ -478,7 +474,7 @@ def build_output(self):
         mean = Dense(self.output_size, activation=MeanAct, kernel_initializer=self.init,
                        kernel_regularizer=l1_l2(self.l1_coef, self.l2_coef),
                        name='mean')(self.decoder_output)
-        output = ColwiseMultLayer([mean, self.sf_layer])
+        output = ColwiseMultLayer()([mean, self.sf_layer])
         output = SliceLayer(0, name='slice')([output, disp, pi])
 
         zinb = ZINB(pi, theta=disp, ridge_lambda=self.ridge, debug=self.debug)
@@ -508,7 +504,7 @@ def build_output(self):
         disp = ConstantDispersionLayer(name='dispersion')
         mean = disp(mean)
 
-        output = ColwiseMultLayer([mean, self.sf_layer])
+        output = ColwiseMultLayer()([mean, self.sf_layer])
 
         zinb = ZINB(pi, theta=disp.theta_exp, ridge_lambda=self.ridge, debug=self.debug)
         self.loss = zinb.loss
@@ -622,7 +618,7 @@ def build(self):
                 # Use separate act. layers to give user the option to get pre-activations
                 # of layers when requested
                 if self.activation in advanced_activations:
-                    last_hidden = keras.layers.__dict__[self.activation](name='%s_act'%layer_name)(last_hidden)
+                    last_hidden = tf.keras.layers.__dict__[self.activation](name='%s_act'%layer_name)(last_hidden)
                 else:
                     last_hidden = Activation(self.activation, name='%s_act'%layer_name)(last_hidden)
 
@@ -646,7 +642,7 @@ def build_output(self):
                        kernel_regularizer=l1_l2(self.l1_coef, self.l2_coef),
                        name='mean')(self.last_hidden_mean)
 
-        output = ColwiseMultLayer([mean, self.sf_layer])
+        output = ColwiseMultLayer()([mean, self.sf_layer])
         output = SliceLayer(0, name='slice')([output, disp, pi])
 
         zinb = ZINB(pi, theta=disp, ridge_lambda=self.ridge, debug=self.debug)
@@ -726,7 +722,7 @@ def build(self):
                 # Use separate act. layers to give user the option to get pre-activations
                 # of layers when requested
                 if self.activation in advanced_activations:
-                    last_hidden = keras.layers.__dict__[self.activation](name='%s_act'%layer_name)(last_hidden)
+                    last_hidden = tf.keras.layers.__dict__[self.activation](name='%s_act'%layer_name)(last_hidden)
                 else:
                     last_hidden = Activation(self.activation, name='%s_act'%layer_name)(last_hidden)
 
@@ -747,7 +743,7 @@ def build_output(self):
                        kernel_regularizer=l1_l2(self.l1_coef, self.l2_coef),
                        name='mean')(self.last_hidden_mean)
 
-        output = ColwiseMultLayer([mean, self.sf_layer])
+        output = ColwiseMultLayer()([mean, self.sf_layer])
         output = SliceLayer(0, name='slice')([output, disp])
 
         nb = NB(theta=disp, debug=self.debug)
diff --git a/dca/test.py b/dca/test.py
index 15f0a9e..246ad83 100644
--- a/dca/test.py
+++ b/dca/test.py
@@ -1,59 +1,73 @@
 import numpy as np
 import scanpy as sc
 
+from unittest import TestCase
+
 from .api import dca
 
-def test_api():
-    adata = sc.datasets.paul15()
-    epochs = 1
-
-    # simple tests for denoise
-    ret = dca(adata, mode='denoise', copy=True, epochs=epochs, verbose=True)
-    assert not np.allclose(ret.X[:10], adata.X[:10])
-
-    ret, _ = dca(adata, mode='denoise', ae_type='nb-conddisp', copy=True, epochs=epochs,
-              return_model=True, return_info=True)
-    assert not np.allclose(ret.X[:10], adata.X[:10])
-    assert 'X_dca_dispersion' in ret.obsm_keys()
-    assert _ is not None
-
-    ret = dca(adata, mode='denoise', ae_type='nb', copy=True, epochs=epochs,
-              return_model=False, return_info=True)
-    assert not np.allclose(ret.X[:10], adata.X[:10])
-    assert 'X_dca_dispersion' in ret.var_keys()
-
-    ret = dca(adata, mode='denoise', ae_type='zinb', copy=True, epochs=epochs,
-              return_model=False, return_info=True)
-    assert not np.allclose(ret.X[:10], adata.X[:10])
-    assert 'X_dca_dropout' in ret.obsm_keys()
-    assert 'dca_loss_history' in ret.uns_keys()
-
-    ret = dca(adata, mode='denoise', ae_type='zinb-elempi', copy=True, epochs=epochs,
-              return_model=False, return_info=True)
-    assert not np.allclose(ret.X[:10], adata.X[:10])
-    assert 'X_dca_dropout' in ret.obsm_keys()
-    assert 'dca_loss_history' in ret.uns_keys()
-
-    ret = dca(adata, mode='denoise', ae_type='zinb-elempi', copy=True, epochs=epochs,
-              return_model=False, return_info=True, network_kwds={'sharedpi': True})
-    assert not np.allclose(ret.X[:10], adata.X[:10])
-    assert 'X_dca_dropout' in ret.obsm_keys()
-    assert 'dca_loss_history' in ret.uns_keys()
-
-    # simple tests for latent
-    hid_size = (10, 2, 10)
-    ret = dca(adata, mode='latent', hidden_size=hid_size, copy=True, epochs=epochs)
-    assert 'X_dca' in ret.obsm_keys()
-    assert ret.obsm['X_dca'].shape[1] == hid_size[1]
-
-    ret = dca(adata, mode='latent', ae_type='nb-conddisp', hidden_size=hid_size, copy=True, epochs=epochs)
-    assert 'X_dca' in ret.obsm_keys()
-    assert ret.obsm['X_dca'].shape[1] == hid_size[1]
-
-    ret = dca(adata, mode='latent', ae_type='nb', hidden_size=hid_size, copy=True, epochs=epochs, return_info=True)
-    assert 'X_dca' in ret.obsm_keys()
-    assert ret.obsm['X_dca'].shape[1] == hid_size[1]
-
-    ret = dca(adata, mode='latent', ae_type='zinb', hidden_size=hid_size, copy=True, epochs=epochs)
-    assert 'X_dca' in ret.obsm_keys()
-    assert ret.obsm['X_dca'].shape[1] == hid_size[1]
+class TestAPI(TestCase):
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls.adata = sc.datasets.paul15()
+        cls.epochs = 1
+        cls.hid_size = (10, 2, 10)
+
+    def test_denoise_simple(self):
+        ret = dca(self.adata, mode='denoise', copy=True, epochs=self.epochs, verbose=True)
+        assert not np.allclose(ret.X[:10], self.adata.X[:10])
+
+    def test_denoise_nb_conddisp(self):
+        ret, _ = dca(self.adata, mode='denoise', ae_type='nb-conddisp', copy=True, epochs=self.epochs,
+                  return_model=True, return_info=True)
+        assert not np.allclose(ret.X[:10], self.adata.X[:10])
+        assert 'X_dca_dispersion' in ret.obsm_keys()
+        assert _ is not None
+
+    def test_denoise_nb(self):
+        ret = dca(self.adata, mode='denoise', ae_type='nb', copy=True, epochs=self.epochs,
+                  return_model=False, return_info=True)
+        assert not np.allclose(ret.X[:10], self.adata.X[:10])
+        assert 'X_dca_dispersion' in ret.var_keys()
+
+    def test_denoise_zinb(self):
+        ret = dca(self.adata, mode='denoise', ae_type='zinb', copy=True, epochs=self.epochs,
+                  return_model=False, return_info=True)
+        assert not np.allclose(ret.X[:10], self.adata.X[:10])
+        assert 'X_dca_dropout' in ret.obsm_keys()
+        assert 'dca_loss_history' in ret.uns_keys()
+
+    def test_denoise_zinb_elempi(self):
+        ret = dca(self.adata, mode='denoise', ae_type='zinb-elempi', copy=True, epochs=self.epochs,
+                  return_model=False, return_info=True)
+        assert not np.allclose(ret.X[:10], self.adata.X[:10])
+        assert 'X_dca_dropout' in ret.obsm_keys()
+        assert 'dca_loss_history' in ret.uns_keys()
+
+    def test_denoise_zinb_elempi_sharedpi(self):
+        ret = dca(self.adata, mode='denoise', ae_type='zinb-elempi', copy=True, epochs=self.epochs,
+                  return_model=False, return_info=True, network_kwds={'sharedpi': True})
+        assert not np.allclose(ret.X[:10], self.adata.X[:10])
+        assert 'X_dca_dropout' in ret.obsm_keys()
+        assert 'dca_loss_history' in ret.uns_keys()
+
+    def test_latent_simple(self):
+        # simple tests for latent
+        ret = dca(self.adata, mode='latent', hidden_size=self.hid_size, copy=True, epochs=self.epochs)
+        assert 'X_dca' in ret.obsm_keys()
+        assert ret.obsm['X_dca'].shape[1] == self.hid_size[1]
+
+    def test_latent_nb_conddisp(self):
+        ret = dca(self.adata, mode='latent', ae_type='nb-conddisp', hidden_size=self.hid_size, copy=True, epochs=self.epochs)
+        assert 'X_dca' in ret.obsm_keys()
+        assert ret.obsm['X_dca'].shape[1] == self.hid_size[1]
+
+    def test_latent_nb(self):
+        ret = dca(self.adata, mode='latent', ae_type='nb', hidden_size=self.hid_size, copy=True, epochs=self.epochs, return_info=True)
+        assert 'X_dca' in ret.obsm_keys()
+        assert ret.obsm['X_dca'].shape[1] == self.hid_size[1]
+
+    def test_latent_zinb(self):
+        ret = dca(self.adata, mode='latent', ae_type='zinb', hidden_size=self.hid_size, copy=True, epochs=self.epochs)
+        assert 'X_dca' in ret.obsm_keys()
+        assert ret.obsm['X_dca'].shape[1] == self.hid_size[1]
diff --git a/dca/train.py b/dca/train.py
index 6171f26..72223c9 100644
--- a/dca/train.py
+++ b/dca/train.py
@@ -22,14 +22,11 @@
 
 from . import io
 from .network import AE_types
-from .hyper import hyper
 
 import numpy as np
 import tensorflow as tf
-import keras.optimizers as opt
-from keras.callbacks import TensorBoard, ModelCheckpoint, EarlyStopping, ReduceLROnPlateau
-from keras import backend as K
-from keras.preprocessing.image import Iterator
+import tensorflow.keras.optimizers as opt
+from tensorflow.keras.callbacks import TensorBoard, ModelCheckpoint, EarlyStopping, ReduceLROnPlateau
 
 
 def train(adata, network, output_dir=None, optimizer='RMSprop', learning_rate=None,
@@ -38,14 +35,8 @@ def train(adata, network, output_dir=None, optimizer='RMSprop', learning_rate=No
           validation_split=0.1, tensorboard=False, verbose=True, threads=None,
           **kwds):
 
-    tf.compat.v1.keras.backend.set_session(
-        tf.compat.v1.Session(
-            config=tf.compat.v1.ConfigProto(
-                intra_op_parallelism_threads=threads,
-                inter_op_parallelism_threads=threads,
-            )
-        )
-    )
+    tf.config.threading.set_inter_op_parallelism_threads(threads)
+    tf.config.threading.set_intra_op_parallelism_threads(threads)
     model = network.model
     loss = network.loss
     if output_dir is not None:
@@ -102,14 +93,9 @@ def train(adata, network, output_dir=None, optimizer='RMSprop', learning_rate=No
 
 def train_with_args(args):
 
-    tf.compat.v1.keras.backend.set_session(
-        tf.compat.v1.Session(
-            config=tf.compat.v1.ConfigProto(
-                intra_op_parallelism_threads=args.threads,
-                inter_op_parallelism_threads=args.threads,
-            )
-        )
-    )
+
+    tf.config.threading.set_inter_op_parallelism_threads(args.threads)
+    tf.config.threading.set_intra_op_parallelism_threads(args.threads)
     # set seed for reproducibility
     random.seed(42)
     np.random.seed(42)
@@ -118,6 +104,7 @@ def train_with_args(args):
 
     # do hyperpar optimization and exit
     if args.hyper:
+        from .hyper import hyper
         hyper(args)
         return
 
diff --git a/dca/utils.py b/dca/utils.py
index 969e3b7..80e51b6 100644
--- a/dca/utils.py
+++ b/dca/utils.py
@@ -1,7 +1,6 @@
 import scanpy as sc
 import matplotlib.pyplot as plt
 import numpy as np
-import pandas as pd
 import seaborn as sns
 import scipy as sp
 import tensorflow as tf
diff --git a/setup.py b/setup.py
index 0c70f64..dedd57b 100644
--- a/setup.py
+++ b/setup.py
@@ -2,14 +2,14 @@
 
 setup(
     name='DCA',
-    version='0.3.3',
+    version='0.3.4',
     description='Count autoencoder for scRNA-seq denoising',
     author='Gokcen Eraslan',
     author_email="gokcen.eraslan@gmail.com",
     packages=['dca'],
     install_requires=['numpy>=1.7',
-                      'keras>=2.4,<2.6',
-                      'tensorflow>=2.0,<2.5',
+                      'tensorflow>=2.0,<2.11,!=2.6',
+                      'protobuf<=3.20',
                       'h5py',
                       'six>=1.10.0',
                       'scikit-learn',
@@ -17,6 +17,7 @@
                       'kopt',
                       'pandas' #for preprocessing
                       ],
+    extras_require={"test":["pytest"]},
     url='https://github.com/theislab/dca',
     entry_points={
         'console_scripts': [