diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
new file mode 100644
index 0000000..cfb6920
--- /dev/null
+++ b/.github/workflows/publish.yml
@@ -0,0 +1,29 @@
+name: Upload Python Package
+
+on:
+  release:
+    types: [created]
+
+jobs:
+
+  deploy:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install setuptools wheel twine
+    - name: Build and publish
+      env:
+        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
+        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
+      run: |
+        python setup.py sdist bdist_wheel
+        twine upload dist/*
\ No newline at end of file
diff --git a/setup.py b/setup.py
index c2612ad..3b02221 100644
--- a/setup.py
+++ b/setup.py
@@ -2,11 +2,11 @@
 from setuptools import find_packages
 
 setup(name='sparktorch',
-      version='0.1.2',
+      version='0.2.0',
       description='Distributed training of PyTorch networks on Apache Spark with ML Pipeline support',
-      keywords = ['pytorch', 'spark', 'sparktorch', 'machine learning', 'deep learning'],
+      keywords=['pytorch', 'spark', 'sparktorch', 'machine learning', 'deep learning'],
       url='https://github.com/dmmiller612/sparktorch',
-      download_url='https://github.com/dmmiller612/sparktorch/archive/0.1.2.tar.gz',
+      download_url='https://github.com/dmmiller612/sparktorch/archive/0.2.0.tar.gz',
       author='Derek Miller',
       author_email='dmmiller612@gmail.com',
       long_description=open("README.md", "r", encoding='utf-8').read(),
diff --git a/sparktorch/distributed.py b/sparktorch/distributed.py
index 3cbaa19..0e0ee0b 100644
--- a/sparktorch/distributed.py
+++ b/sparktorch/distributed.py
@@ -114,7 +114,7 @@ def handle_model(
     # Loaded the model
     model = torch_obj.model.to(device)
 
-    if compile is not None:
+    if compile_mode is not None:
         torch.compile(compile_mode)
 
     model.train()
@@ -178,7 +178,7 @@ def handle_model(
         # Distributed part of training.
         for param in model.parameters():
             dist.all_reduce(param.grad.data, op=torch.distributed.ReduceOp.SUM)
-            param.grad.data /= (world_size - 1)
+            param.grad.data /= world_size
 
         # Processes the early stop work
         loss_distributed = None
@@ -186,7 +186,7 @@ def handle_model(
             loss_to_use = val_loss if val_loss is not None else loss
 
             dist.all_reduce(loss_to_use, op=torch.distributed.ReduceOp.SUM)
-            loss_distributed = loss_to_use.item() / (world_size - 1)
+            loss_distributed = loss_to_use.item() / world_size
             stop = es.step(loss_distributed)
             if stop:
                 should_stop = should_stop + 1.0
@@ -223,7 +223,7 @@ def train_distributed(
     :param rdd: The rdd of data to run on the model.
     :param torch_obj: The torch object as a string that includes the model and param shapes.
     :param iters: Number of iterations for training.
-    :param partition_shuffles: Number of partition shuffles (Need to implement).
+    :param partition_shuffles: Number of partition shuffles.
     :param verbose: Verbosity of logs.
     :param mini_batch: Mini batch for each iteration of training.
     :param validation_pct: How many items to validate.