diff --git a/flashinfer/autotuner.py b/flashinfer/autotuner.py index 6b6a0c3e48..9f5fb67489 100644 --- a/flashinfer/autotuner.py +++ b/flashinfer/autotuner.py @@ -61,9 +61,9 @@ def __post_init__(self): # Set default tensor_initializers if not provided if self.tensor_initializers is None: self.tensor_initializers = [ - lambda shapes, dtype, device: torch.randn(shapes, device=device).to( - dtype - ) + lambda shapes, dtype, device: ( + torch.rand(shapes, device=device) * 10 - 5 + ).to(dtype) for _ in range(len(self.input_idx)) ] @@ -761,8 +761,8 @@ def _create_tensor_like( def _prepare_input_tensors( self, profile: OptimizationProfile, inputs: List[torch.Tensor] ) -> List[torch.Tensor]: - default_initializer = lambda shapes, dtype, device: torch.rand( - shapes, device=device + default_initializer = lambda shapes, dtype, device: ( + torch.rand(shapes, device=device) * 10 - 5 ).to(dtype) tensors = [] for i, p in enumerate(profile.shapes):