send tensors to the correct device when loading from safetensors file with...

def read_state_dict(checkpoint_file, print_global_state=False, map_location=None):

    _, extension = os.path.splitext(checkpoint_file)

    if extension.lower() == ".safetensors":

        if not shared.opts.disable_mmap_load_safetensors:

        device = map_location or shared.weight_load_location or devices.get_optimal_device_name()

        if not shared.opts.disable_mmap_load_safetensors:

            pl_sd = safetensors.torch.load_file(checkpoint_file, device=device)

        else:

            pl_sd = safetensors.torch.load(open(checkpoint_file, 'rb').read())

            pl_sd = {k: v.to(device) for k, v in pl_sd.items()}

    else:

        pl_sd = torch.load(checkpoint_file, map_location=map_location or shared.weight_load_location)

    "multiple_tqdm": OptionInfo(True, "Add a second progress bar to the console that shows progress for an entire job."),

    "print_hypernet_extra": OptionInfo(False, "Print extra hypernetwork information to console."),

    "list_hidden_files": OptionInfo(True, "Load models/files in hidden directories").info("directory is hidden if its name starts with \".\""),

    "disable_mmap_load_safetensors": OptionInfo(False, "Disable memmapping for loading .safetensors files (fixes very slow loading speed in some cases)."),

    "disable_mmap_load_safetensors": OptionInfo(False, "Disable memmapping for loading .safetensors files.").info("fixes very slow loading speed in some cases"),

}))

options_templates.update(options_section(('training', "Training"), {