mirror of
https://github.com/huggingface/text-generation-inference.git
synced 2025-09-09 11:24:53 +00:00
use hpu set seed
Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
This commit is contained in:
parent
0e45bca624
commit
babba0a5e5
@ -552,6 +552,11 @@ def pad_next_token_chooser_parameters(
|
|||||||
|
|
||||||
class Sampling:
|
class Sampling:
|
||||||
def __init__(self, seed: int, device: str = "cpu"):
|
def __init__(self, seed: int, device: str = "cpu"):
|
||||||
|
if device in ["hpu", torch.device("hpu")]:
|
||||||
|
import habana_frameworks.torch.hpu.random as htrandom
|
||||||
|
|
||||||
|
self.generator = htrandom.default_generators[0].manual_seed(seed)
|
||||||
|
else:
|
||||||
self.generator = torch.Generator("cpu")
|
self.generator = torch.Generator("cpu")
|
||||||
self.generator.manual_seed(seed)
|
self.generator.manual_seed(seed)
|
||||||
self.seed = seed
|
self.seed = seed
|
||||||
|
Loading…
Reference in New Issue
Block a user