Suppresses pytorch warning

#5
Files changed (1) hide show
  1. modeling_hf_nomic_bert.py +1 -1
modeling_hf_nomic_bert.py CHANGED
@@ -75,7 +75,7 @@ def state_dict_from_pretrained(model_name, safe_serialization=False, device=None
75
  if load_safe:
76
  loader = partial(safe_load_file, device=mapped_device)
77
  else:
78
- loader = partial(torch.load, map_location=mapped_device)
79
 
80
  if is_sharded:
81
  # resolved_archive_file becomes a list of files that point to the different
 
75
  if load_safe:
76
  loader = partial(safe_load_file, device=mapped_device)
77
  else:
78
+ loader = partial(torch.load, map_location=mapped_device, weights_only=True)
79
 
80
  if is_sharded:
81
  # resolved_archive_file becomes a list of files that point to the different