get_model() — pytorch Function Reference

Architecture documentation for the get_model() function in dist_util.py from the pytorch codebase.

Function python DynamoBenchmarks Configuration called by 1

Entity Profile

DynamoBenchmarks→ Configuration→ get_model() — pytorch Function Reference

Dependency Diagram

graph TD
  3ea941c9_72e7_2829_0157_0cee5513e057["get_model()"]
  28abd701_daff_7831_6354_47eeb9f5e68d["parser()"]
  28abd701_daff_7831_6354_47eeb9f5e68d -->|calls| 3ea941c9_72e7_2829_0157_0cee5513e057
  style 3ea941c9_72e7_2829_0157_0cee5513e057 fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/dist_util.py lines 88–105

def get_model(args):
    if args.torchbench_model:
        setup_torchbench_cwd()
        module = importlib.import_module(
            f"torchbenchmark.models.{args.torchbench_model}"
        )
        benchmark_cls = getattr(module, "Model", None)
        bm = benchmark_cls(test="train", device=args.device, batch_size=args.batch_size)
        model, inputs = bm.get_module()
    elif args.toy_model:
        model = ToyModel()
        inputs = (torch.randn(20, 10),)
    else:
        raise argparse.ArgumentError(
            args.torchbench_model, message="Must specify a model"
        )

    return model, inputs