metaopt · XuehaiPan · Sep 9, 2022 · Sep 9, 2022 · Sep 9, 2022 · Sep 9, 2022
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -58,6 +58,14 @@ jobs:
           python-version: "3.7 - 3.10"
           update-environment: true
 
+      - name: Set __release__
+        if: startsWith(github.ref, 'refs/tags/') || github.event_name == 'workflow_dispatch'
+        run: |
+          sed -i -E 's/^__release__\s*=.*$/__release__ = True/' torchopt/version.py
+
+      - name: Print version
+        run: python setup.py --version
+
       - name: Install dependencies
         run: python -m pip install --upgrade pip setuptools wheel build
 
@@ -83,6 +91,21 @@ jobs:
           submodules: "recursive"
           fetch-depth: 0
 
+      - name: Set up Python
+        id: py
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.7 - 3.10"
+          update-environment: true
+
+      - name: Set __release__
+        if: startsWith(github.ref, 'refs/tags/') || github.event_name == 'workflow_dispatch'
+        run: |
+          sed -i -E 's/^__release__\s*=.*$/__release__ = True/' torchopt/version.py
+
+      - name: Print version
+        run: python setup.py --version
+
       - name: Build wheels
         uses: pypa/cibuildwheel@v2.8.1
         with:

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -21,7 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
-
+- Fix the step count tensor (`shape=(1,)`) can change the shape of the scalar updates (`shape=()`) by [@XuehaiPan](https://github.com/XuehaiPan) in [#71](https://github.com/metaopt/torchopt/pull/71).
 
 ### Removed
 

diff --git a/torchopt/_src/transform.py b/torchopt/_src/transform.py
@@ -161,7 +161,7 @@ def _scale_by_schedule(
 
     def init_fn(params):
         zero = tree_map(  # count init
-            lambda t: torch.zeros(1, dtype=torch.int32, device=t.device), params
+            lambda t: torch.zeros(1, dtype=torch.int32, device=t.device).squeeze_(), params
         )
         return ScaleByScheduleState(count=zero)
 
@@ -299,7 +299,7 @@ def _scale_by_adam(
 
     def init_fn(params):
         zero = tree_map(  # count init
-            lambda t: torch.zeros(1, dtype=torch.int32, device=t.device), params
+            lambda t: torch.zeros(1, dtype=torch.int32, device=t.device).squeeze_(), params
         )
         mu = tree_map(  # first moment
             lambda t: torch.zeros_like(t, requires_grad=moment_requires_grad), params
@@ -426,7 +426,7 @@ def update_fn(updates, state, *, params=None, inplace=True):
 
     def init_fn(params):
         zero = tree_map(  # count init
-            lambda t: torch.zeros(1, dtype=torch.int32, device=t.device), params
+            lambda t: torch.zeros(1, dtype=torch.int32, device=t.device).squeeze_(), params
         )
         mu = tree_map(  # first moment
             lambda t: torch.zeros_like(t, requires_grad=moment_requires_grad), params

diff --git a/torchopt/version.py b/torchopt/version.py
@@ -15,3 +15,37 @@
 """TorchOpt: a high-performance optimizer library built upon PyTorch."""
 
 __version__ = '0.5.0'
+__license__ = 'Apache License, Version 2.0'
+__author__ = 'TorchOpt Contributors'
+__release__ = False
+
+if not __release__:
+    import os
+    import subprocess
+
+    try:
+        prefix, sep, suffix = (
+            subprocess.check_output(
+                ['git', 'describe', '--abbrev=7'],
+                cwd=os.path.dirname(os.path.abspath(__file__)),
+                stderr=subprocess.DEVNULL,
+                universal_newlines=True,
+            )
+            .strip()
+            .lstrip('v')
+            .replace('-', '.dev', 1)
+            .replace('-', '+', 1)
+            .partition('.dev')
+        )
+        if sep:
+            version_prefix, dot, version_tail = prefix.rpartition('.')
+            prefix = f'{version_prefix}{dot}{int(version_tail) + 1}'
+            __version__ = sep.join((prefix, suffix))
+            del version_prefix, dot, version_tail
+        else:
+            __version__ = prefix
+        del prefix, sep, suffix
+    except (OSError, subprocess.CalledProcessError):
+        pass
+
+    del os, subprocess
Original file line number	Diff line number	Diff line change
Expand Up		@@ -21,7 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

		### Fixed


		- Fix the step count tensor (`shape=(1,)`) can change the shape of the scalar updates (`shape=()`) by [@XuehaiPan](https://github.com/XuehaiPan) in [#71](https://github.com/metaopt/torchopt/pull/71).

		### Removed

Expand Down