diff --git a/.github/actions/exp/jl-dvc/action.yml b/.github/actions/exp/jl-dvc/action.yml
index ceaacfa..595601f 100644
--- a/.github/actions/exp/jl-dvc/action.yml
+++ b/.github/actions/exp/jl-dvc/action.yml
@@ -17,14 +17,13 @@ runs:
   using: "composite"
   steps:      
       - name: prepare script
-        if: steps.check.outputs.triggered == 'true'      
         env:
-          CUR_BRANCH: ${{ steps.pr_data.outputs.branch }}
-          CUR_PR_ID: ${{ steps.pr_data.outputs.number }}
-          GDRIVE_CREDENTIAL: ${{ secrets.GDRIVE_CREDENTIAL }}
-          GH_ACCESS_TOKEN: ${{ secrets.GH_ACCESS_TOKEN }}
-          CLOUD_ID: ${{ env.JARVISLABS_ID }}
-          CLOUD_AT: ${{ env.JARVISLABS_ACCESS_TOKEN }}
+          CUR_BRANCH: ${{ inputs.CUR_BRANCH }}
+          CUR_PR_ID: ${{ inputs.CUR_PR_ID }}
+          GDRIVE_CREDENTIAL: ${{ inputs.GDRIVE_CREDENTIAL }}
+          GH_ACCESS_TOKEN: ${{ inputs.GH_ACCESS_TOKEN }}
+          CLOUD_ID: ${{ inputs.JARVISLABS_ID }}
+          CLOUD_AT: ${{ inputs.JARVISLABS_ACCESS_TOKEN }}
         run: |
           envsubst \
             '$CUR_BRANCH, \
@@ -36,32 +35,28 @@ runs:
              < scripts/experiments.sh \
              > scripts/experiments_tmp.sh
              
-      - name: install jarvislabs client
-        if: steps.check.outputs.triggered == 'true'      
+      - name: install jarvislabs client 
         run: |
           pip install typer
           pip install git+https://github.com/jarvislabsai/jlclient.git
           
       - name: add script to jarvislabs
         id: add_script
-        if: steps.check.outputs.triggered == 'true'        
         run: |
           python clouds/jarvislabs.py \
             script add \
-              ${{ env.CLOUD_AT }} ${{ env.CLOUD_ID }} scripts/experiments_tmp.sh \
+              ${{ inputs.JARVISLABS_ACCESS_TOKEN }} ${{ inputs.JARVISLABS_ID }} scripts/experiments_tmp.sh \
                 > outputs.txt
           echo "::set-output name=script_id::$(cat outputs.txt)"
           
       - name: create vm on jarvislabs
-        if: steps.check.outputs.triggered == 'true'      
         run: |
           python clouds/jarvislabs.py \
             vm create \
-              ${{ env.CLOUD_AT }} ${{ env.CLOUD_ID }} ${{ steps.add_script.outputs.script_id }}
+              ${{ inputs.JARVISLABS_ACCESS_TOKEN }} ${{ inputs.JARVISLABS_ID }} ${{ steps.add_script.outputs.script_id }}
               
       - name: remove script from jarvislabs
-        if: steps.check.outputs.triggered == 'true'      
         run: |
           python clouds/jarvislabs.py \
             script remove \
-              ${{ env.CLOUD_AT }} ${{ env.CLOUD_ID }} ${{ steps.add_script.outputs.script_id }}
\ No newline at end of file
+              ${{ inputs.JARVISLABS_ACCESS_TOKEN }} ${{ inputs.JARVISLABS_ID }} ${{ steps.add_script.outputs.script_id }}
\ No newline at end of file
diff --git a/.github/actions/exp/jl-wandb/action.yml b/.github/actions/exp/jl-wandb/action.yml
index 4296546..5cf12e1 100644
--- a/.github/actions/exp/jl-wandb/action.yml
+++ b/.github/actions/exp/jl-wandb/action.yml
@@ -22,14 +22,14 @@ runs:
   steps:
       - name: prepare script
         env:
-          CUR_BRANCH: $CUR_BRANCH
-          CUR_PR_ID: $CUR_PR_ID
-          GDRIVE_CREDENTIAL: $GDRIVE_CREDENTIAL
-          GH_ACCESS_TOKEN: $GH_ACCESS_TOKEN
-          CLOUD_ID: $JARVISLABS_ID
-          CLOUD_AT: $JARVISLABS_ACCESS_TOKEN
-          WANDB_PROJECT: $WANDB_PROJECT
-          WANDB_API_KEY: $WANDB_API_KEY
+          CUR_BRANCH: ${{ inputs.CUR_BRANCH }}
+          CUR_PR_ID: ${{ inputs.CUR_PR_ID }}
+          GDRIVE_CREDENTIAL: ${{ inputs.GDRIVE_CREDENTIAL }}
+          GH_ACCESS_TOKEN: ${{ inputs.GH_ACCESS_TOKEN }}
+          CLOUD_ID: ${{ inputs.JARVISLABS_ID }}
+          CLOUD_AT: ${{ inputs.JARVISLABS_ACCESS_TOKEN }}
+          WANDB_PROJECT: ${{ inputs.WANDB_PROJECT }}
+          WANDB_API_KEY: ${{ inputs.WANDB_API_KEY }}
         run: |
           envsubst \
             '$CUR_BRANCH, \
@@ -37,32 +37,39 @@ runs:
              $GDRIVE_CREDENTIAL, \
              $GH_ACCESS_TOKEN, \
              $CLOUD_ID, \
-             $CLOUD_AT' \
-             < scripts/js_exp_wandb.sh \
-             > scripts/js_exp_wandb_tmp.sh
+             $CLOUD_AT, \
+             $WANDB_PROJECT, \
+             $WANDB_API_KEY' \
+             < scripts/jl_exp_wandb.sh \
+             > scripts/jl_exp_wandb_tmp.sh
+        shell: bash             
              
       - name: install jarvislabs client
         run: |
           pip install typer
           pip install git+https://github.com/jarvislabsai/jlclient.git
+        shell: bash          
           
       - name: add script to jarvislabs
         id: add_script
         run: |
           python clouds/jarvislabs.py \
             script add \
-              ${{ env.CLOUD_AT }} ${{ env.CLOUD_ID }} scripts/js_exp_wandb_tmp.sh \
+              ${{ inputs.JARVISLABS_ACCESS_TOKEN }} ${{ inputs.JARVISLABS_ID }} scripts/jl_exp_wandb_tmp.sh \
                 > outputs.txt
           echo "::set-output name=script_id::$(cat outputs.txt)"
+        shell: bash          
           
       - name: create vm on jarvislabs
         run: |
           python clouds/jarvislabs.py \
             vm create \
-              ${{ env.CLOUD_AT }} ${{ env.CLOUD_ID }} ${{ steps.add_script.outputs.script_id }}
+              ${{ inputs.JARVISLABS_ACCESS_TOKEN }} ${{ inputs.JARVISLABS_ID }} ${{ steps.add_script.outputs.script_id }}
+        shell: bash              
               
       - name: remove script from jarvislabs
         run: |
           python clouds/jarvislabs.py \
             script remove \
-              ${{ env.CLOUD_AT }} ${{ env.CLOUD_ID }} ${{ steps.add_script.outputs.script_id }}
\ No newline at end of file
+              ${{ inputs.JARVISLABS_ACCESS_TOKEN }} ${{ inputs.JARVISLABS_ID }} ${{ steps.add_script.outputs.script_id }}
+        shell: bash              
diff --git a/README.md b/README.md
index a2d67da..c4c9f8f 100644
--- a/README.md
+++ b/README.md
@@ -40,9 +40,9 @@ $ dvc stage add -n train \
 # if you want to use W&B for tracking training progress
 $ dvc stage add -n train \
                 -p train.train_size,train.batch_size,train.epoch,train.lr \
-                -d pipeline/modeling.py -d pipeline/train.py -d data \
+                -d pipeline/modeling.py -d pipeline/train_wandb.py -d data \
                 -o outputs/model \
-                python pipeline/train.py outputs/model
+                python pipeline/train_wandb.py outputs/model
 ```
 11. Run the following `dvc stage add` for evaluate stage
 ```bash
@@ -68,12 +68,13 @@ $ dvc stage add -n evaluate \
 18. Run `git add . && git commit -m "add initial pipeline run" && git push origin main`
 19. Add access token and user email of [JarvisLabs.ai](https://jarvislabs.ai/) to GitHub Secret as `JARVISLABS_ACCESS_TOKEN` and `JARVISLABS_USER_EMAIL`
 20. Add GitHub access token to GitHub Secret as `GH_ACCESS_TOKEN`
-21. Create a PR and write `#train` as in comment (you have to be the onwer of the repo)
+21. Create a PR and write `#train --with dvc` as in comment (you have to be the onwer of the repo)
 
 ### W&B Integration Setup
 
 1. Add W&B's project name to GitHub Secret as `WANDB_PROJECT`
 2. Add W&B's API KEY to GitHub Secret as `WANDB_API_KEY`
+3. Use `#train --with wandb` instead of `#train --with dvc`
 
 ### HuggingFace Integration Setup
 
diff --git a/requirements.txt b/requirements.txt
index 7ca652f..cb329fe 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,5 @@
 dvc[gdrive]==2.10.2
+wandb==0.12.19
 tensorflow==2.8
 typer==0.4.1
 docopt==0.6.2