fix docs string

This commit is contained in:
Kashif Rasul
2026-03-11 17:38:09 +01:00
parent 0134e55a6e
commit d2cb4849b6
+4 -4
View File
@@ -22,7 +22,7 @@ from typing import Optional, Sequence, Union
import numpy as np import numpy as np
import torch import torch
from huggingface_hub import PytorchModelHubMixin, hf_hub_download from huggingface_hub import PyTorchModelHubMixin, hf_hub_download
from safetensors.torch import load_file, save_file from safetensors.torch import load_file, save_file
from torch import nn from torch import nn
@@ -265,7 +265,7 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
class TimesFM_2p5_200M_torch( class TimesFM_2p5_200M_torch(
timesfm_2p5_base.TimesFM_2p5, timesfm_2p5_base.TimesFM_2p5,
PytorchModelHubMixin, PyTorchModelHubMixin,
library_name="timesfm", library_name="timesfm",
repo_url="https://github.com/google-research/timesfm", repo_url="https://github.com/google-research/timesfm",
paper_url="https://arxiv.org/abs/2310.10688", paper_url="https://arxiv.org/abs/2310.10688",
@@ -305,7 +305,7 @@ class TimesFM_2p5_200M_torch(
""" """
Loads a PyTorch safetensors TimesFM model from a local path or the Hugging Loads a PyTorch safetensors TimesFM model from a local path or the Hugging
Face Hub. This method is the backend for the `from_pretrained` class Face Hub. This method is the backend for the `from_pretrained` class
method provided by `ModelHubMixin`. method provided by `PyTorchModelHubMixin`.
""" """
# Determine the path to the model weights. # Determine the path to the model weights.
model_file_path = "" model_file_path = ""
@@ -341,7 +341,7 @@ class TimesFM_2p5_200M_torch(
def _save_pretrained(self, save_directory: Union[str, Path]): def _save_pretrained(self, save_directory: Union[str, Path]):
""" """
Saves the model's state dictionary to a safetensors file. This method Saves the model's state dictionary to a safetensors file. This method
is called by the `save_pretrained` method from `ModelHubMixin`. is called by the `save_pretrained` method from `PyTorchModelHubMixin`.
""" """
if not os.path.exists(save_directory): if not os.path.exists(save_directory):
os.makedirs(save_directory) os.makedirs(save_directory)