From 063108e905dbc7a36964ec1e0cdb43d888468126 Mon Sep 17 00:00:00 2001 From: justheuristic Date: Tue, 28 Mar 2023 23:51:58 +0300 Subject: [PATCH] Update block.py --- src/petals/bloom/block.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/petals/bloom/block.py b/src/petals/bloom/block.py index d420384..868fe96 100644 --- a/src/petals/bloom/block.py +++ b/src/petals/bloom/block.py @@ -5,13 +5,15 @@ See commit history for authorship. """ import os from typing import Optional, Tuple +from packaging import version import torch.nn.quantized.dynamic.modules.linear import transformers from transformers.models.bloom.modeling_bloom import BloomBlock, _expand_mask, _make_causal_mask, build_alibi_tensor if not os.getenv("PETALS_IGNORE_DEPENDENCY_VERSION"): - assert transformers.__version__.startswith("4.26."), "Please install transformers 4.26.1" + assert version.parse("4.26.0") < version.parse(transformers.__version__) < version.parse("5.0.0"), \ + "Please install transformers >=4.26.0,<5.0.0" class WrappedBloomBlock(BloomBlock):