diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 96edc5d7..9cfe674e 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -14,7 +14,6 @@ jobs:
     - name: Build package
       run: pip install wheel && python setup.py sdist bdist_wheel && ls -l dist
     - name: Publish package to TestPyPI
-      if: "github.event.release.prerelease"
       uses: pypa/gh-action-pypi-publish@master
       with:
         user: __token__
diff --git a/README.md b/README.md
index 797a4ab5..e82ade1c 100644
--- a/README.md
+++ b/README.md
@@ -100,7 +100,7 @@ TileDB-ML can be installed:
 
       pip install git+https://github.com/TileDB-Inc/TileDB-ML.git@master
 
-- Soon from PyPi:
+- from PyPi:
 
 [comment]: <> (TileDB-ML is available from either [PyPI]&#40;https://test.pypi.org/project/tiledb-ml/0.1.2.2/&#41; with ``pip``:)
 
diff --git a/setup.cfg b/setup.cfg
index b3ec9e52..2af262c7 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
 name = tiledb-ml
-version = 0.2.1
+version = 0.2.2
 description = Package supports all machine learning functionality for TileDB Embedded and TileDB Cloud
 author = TileDB, Inc.
 author_email = help@tiledb.io
diff --git a/tiledb/ml/models/base.py b/tiledb/ml/models/base.py
index a0156865..550d9970 100644
--- a/tiledb/ml/models/base.py
+++ b/tiledb/ml/models/base.py
@@ -6,13 +6,9 @@
 import json
 import platform
 
-from typing import Optional, Union
+from typing import Optional
 from enum import Enum, unique
 
-from torch.nn import Module
-from tensorflow.python.keras.models import Model
-from sklearn.base import BaseEstimator
-
 
 @unique
 class ModelFileProperties(Enum):
@@ -41,7 +37,7 @@ def __init__(
         uri: str,
         namespace: str = None,
         ctx: tiledb.Ctx = None,
-        model: Union[Module, Model, BaseEstimator] = None,
+        model=None,
     ):
         """
         Base class for saving machine learning models as TileDB arrays