diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index 8158e80..3b7e251 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -16,15 +16,16 @@ on:
 jobs:
   build:
 
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.config.os }}
+
     strategy:
       matrix:
         config:
-          - {python-version: '3.6', doubleml-version: 'release'}
-          - {python-version: '3.7', doubleml-version: 'release'}
-          - {python-version: '3.8', doubleml-version: 'release'}
-          - {python-version: '3.8', doubleml-version: 'dev'}
-          - {python-version: '3.9', doubleml-version: 'release'}
+          - {os: 'ubuntu-20.04', python-version: '3.6', doubleml-version: 'release'}
+          - {os: 'ubuntu-20.04', python-version: '3.7', doubleml-version: 'release'}
+          - {os: 'ubuntu-latest', python-version: '3.8', doubleml-version: 'release'}
+          - {os: 'ubuntu-latest', python-version: '3.8', doubleml-version: 'dev'}
+          - {os: 'ubuntu-latest', python-version: '3.9', doubleml-version: 'release'}
 
     steps:
     - uses: actions/checkout@v3
diff --git a/doubleml_serverless/double_ml_data_aws.py b/doubleml_serverless/double_ml_data_aws.py
index 105e16a..00c68a0 100644
--- a/doubleml_serverless/double_ml_data_aws.py
+++ b/doubleml_serverless/double_ml_data_aws.py
@@ -14,12 +14,12 @@ def __init__(self,
                  x_cols=None,
                  z_cols=None,
                  use_other_treat_as_covariate=True):
-        super().__init__(data,
-                         y_col,
-                         d_cols,
-                         x_cols,
-                         z_cols,
-                         use_other_treat_as_covariate)
+        super().__init__(data=data,
+                         y_col=y_col,
+                         d_cols=d_cols,
+                         x_cols=x_cols,
+                         z_cols=z_cols,
+                         use_other_treat_as_covariate=use_other_treat_as_covariate)
         self._bucket = bucket
         self._file_ending = os.path.splitext(file_key)[1]
         assert self._file_ending in ['.csv']
@@ -80,12 +80,12 @@ def __init__(self,
                  x_cols=None,
                  z_cols=None,
                  use_other_treat_as_covariate=True):
-        super().__init__(data,
-                         y_col,
-                         d_cols,
-                         x_cols,
-                         z_cols,
-                         use_other_treat_as_covariate)
+        super().__init__(data=data,
+                         y_col=y_col,
+                         d_cols=d_cols,
+                         x_cols=x_cols,
+                         z_cols=z_cols,
+                         use_other_treat_as_covariate=use_other_treat_as_covariate)
         self._data_json = data.to_json(orient="columns")
 
     @property