diff --git a/lmdeploy/pytorch/backends/dlinfer/linear.py b/lmdeploy/pytorch/backends/dlinfer/linear.py index f2450a974..ccfb06649 100644 --- a/lmdeploy/pytorch/backends/dlinfer/linear.py +++ b/lmdeploy/pytorch/backends/dlinfer/linear.py @@ -16,7 +16,7 @@ def update_weights(self, weight: torch.Tensor, bias: Optional[torch.Tensor] = None): """update weights.""" - if os.getenv('TORCH_MACA_NN_LAYOUT', 'False').lower() == 'true': + if os.getenv('DLINER_LINEAR_USE_NN_LAYOUT', 'False').lower() == 'true': weight = weight.data.t().contiguous() return weight, bias