Should use weights_only with torch.load

johnma2006 / mamba-minimal

Simple, minimal implementation of the Mamba SSM in one file of PyTorch.

Apache License 2.0

2.62k stars 191 forks source link

Should use weights_only with torch.load #1

Closed kit1980 closed 10 months ago

kit1980 commented 10 months ago

torch.load without weights_only is a potential security issue: https://github.com/johnma2006/mamba-minimal/blob/a2a47e0efcd44902d78efa022e1c81504037f0a5/model.py#L117

Found using https://github.com/pytorch-labs/torchfix:

model.py:117:20: TOR102 [*] torch.load without weights_only parameter is unsafe. Explicitly set weights_only to False only if you trust the data you load and full pickle functionality is needed, otherwise set weights_only=True.

johnma2006 commented 10 months ago

Thanks. I changed it: https://github.com/johnma2006/mamba-minimal/commit/82efa90919c3b5066674216f3edcebb3414a7b8f