Skip to content

Commit

Permalink
Enable torch.load weights_only=True
Browse files Browse the repository at this point in the history
Change torch.load to use the safer option of only loading weights.
The torch.load function is hard coded to use the default value
within the open_clip package. The code will set
TORCH_FORCE_WEIGHTS_ONLY_LOAD environment variable to 'true' if
not already set to override the default.

Fixes #32
  • Loading branch information
johnbradley committed Aug 27, 2024
1 parent 3f65a47 commit b9498fc
Showing 1 changed file with 8 additions and 0 deletions.
8 changes: 8 additions & 0 deletions src/bioclip/predict.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import os
import json
import torch
from torchvision import transforms
Expand All @@ -19,6 +20,13 @@
PRED_CLASSICATION_KEY = "classification"
PRED_SCORE_KEY = "score"


# Use secure model loading unless explicitly disabled.
# See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more.
if not "TORCH_FORCE_WEIGHTS_ONLY_LOAD" in os.environ:
os.environ["TORCH_FORCE_WEIGHTS_ONLY_LOAD"] = "true"


OPENA_AI_IMAGENET_TEMPLATE = [
lambda c: f"a bad photo of a {c}.",
lambda c: f"a photo of many {c}.",
Expand Down

0 comments on commit b9498fc

Please sign in to comment.