fix: poetry file and readme

README.md

@@ -30,22 +30,9 @@ poetry install
 
 ### Define Your Computation
 
-User computation must be defined as **a function** using ZKStats operations and PyTorch functions. The function signature must be `Callable[[State, dict[str, torch.Tensor]], torch.Tensor]`:
+User computation must be defined as **a function** using ZKStats operations and PyTorch functions. The function signature must be `Callable[[State, Args], torch.Tensor]`:
-
-```python
-import torch
-
-from zkstats.computation import State, Args
-
-# User-defined computation
-def user_computation(state: State, args: Args) -> torch.Tensor:
-    # Define your computation here
-    ...
-
-```
-
 - first argument is a `State` object, which contains the statistical functions that ZKStats supports.
-- second argument is a list of PyTorch tensors, the input data. `data[0]` is the first column, `data[1]` is the second column, and so on.
+- second argument is a `Args` object, which is a dictionary of PyTorch tensors, the input data. `Args['column1']` is the first column, `Args['column2']` is the second column, and so on.
 
 For example, we have two columns of data and we want to compute the mean of the medians of the two columns:
 

pyproject.toml

@@ -7,7 +7,8 @@ authors = ["Jern Kunpittaya", "Kevin Chia"]
 [tool.poetry.dependencies]
 python = "^3.9"
 ezkl = "9.1.0"
-torch = "^2.1.1"
+# fix torch version to 2.2.0 due to a weird issue when upgrading to 2.4.1
+torch = "2.2.0"
 requests = "^2.31.0"
 scipy = "^1.11.4"
 numpy = "^1.26.2"
@@ -15,7 +16,6 @@ matplotlib = "^3.8.2"
 statistics = "^1.0.3"
 onnx = "^1.15.0"
 click = "^8.1.7"
-six = "^1.16.0"
 
 [tool.poetry.dev-dependencies]
 ipykernel = "^6.27.1"