class_docstr: An operation which saves an OpenAlphaTensor model.
  The model parameters are stored in a json file, while the model weights
  are stored in a .pt file.
class_docstr_tok: ["An","operation","which","saves","an","OpenAlphaTensor","model",".","The","model","parameters","are","stored","in","a","json","file",",","while","the","model","weights","are","stored","in","a",".pt","file","."]
docstr: Saves the OpenAlphaTensor model.
  
  Args:
      model (AlphaTensorModel): OpenAlphaTensor model to be saved.
      save_dir (str): Directory where the model will be saved.
docstr_tok: ["Saves","the","OpenAlphaTensor","model",".","Args",":","model","(","AlphaTensorModel",")",":","OpenAlphaTensor","model","to","be","saved",".","save_dir","(","str",")",":","Directory","where","the","model","will","be","saved","."]
code: def execute(
      self,
      model: AlphaTensorModel,
      save_dir: str,
  ):
      """Saves the OpenAlphaTensor model.
  
      Args:
          model (AlphaTensorModel): OpenAlphaTensor model to be saved.
          save_dir (str): Directory where the model will be saved.
      """
      save_dir = Path(save_dir if save_dir else ".")
      save_dir.mkdir(parents=True, exist_ok=True)
      torch.save(model.state_dict(), save_dir / "final_model.pt")
      model_params = {
          "input_size": model.input_size,
          "tensor_length": model.tensor_length,
          "scalars_size": 1,
          "emb_dim": model.emb_dim,
          "n_steps": model.n_steps,
          "n_logits": model.n_logits,
          "n_samples": model.n_samples,
      }
      # save parameters in a json file
      with open(save_dir / "model_params.json", "w") as f:
          json.dump(model_params, f)
  
imports: import json
  from pathlib import Path
  from typing import Any
  import torch
  from nebullvm.operations.base import Operation
  from open_alpha_tensor.core.modules.alpha_tensor import AlphaTensorModel
label_desc: Use image node_id 2 for calling the SaveModelOp obj's underlying member method code with example usage: obj.execute(model, save_dir) without return types