Merge 0.0.8.post1 into main (#122)

* Fix cuda context bug caused by PreTrainedModel (#113)

* Fix cuda context bug caused by PreTrainedModel

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

* Flake8 style fixes

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

---------

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

* update version (#115)

---------

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

Author: VibhuJawa

crossfit/backend/torch/hf/memory_curve_utils.py

@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
 
 import joblib
 import numpy as np
 import torch
+import transformers
 from sklearn.linear_model import LinearRegression
 from tqdm import tqdm
-from transformers import PreTrainedModel
 
 from crossfit.utils.model_adapter import adapt_model_input
 from crossfit.utils.torch_utils import (
@@ -29,7 +30,7 @@
 
 
 def fit_memory_estimate_curve(
-    model: PreTrainedModel,
+    model: "transformers.PreTrainedModel",
     path_or_name: str,
     start_batch_size: int = 1,
     end_batch_size: int = 2048,

examples/dask_aggregate_bench.py

@@ -53,11 +53,7 @@
     columns = [f"I{i}" for i in range(1, ncolumns + 1)]
     if groupby:
         columns += groupby if isinstance(groupby, list) else [groupby]
-    ddf = dd.read_parquet(
-        path,
-        blocksize=blocksize,
-        columns=columns,
-    )
+    ddf = dd.read_parquet(path, blocksize=blocksize, columns=columns)
     print(f"\nddf: {ddf}\n")
 
     # Aggregate moments (mean, var, std)

setup.py

@@ -18,7 +18,7 @@
 
 from setuptools import find_packages, setup
 
-VERSION = "0.0.8"
+VERSION = "0.0.8.post1"
 
 
 def get_long_description():