from __future__ import annotations

import pytest
import torch

from jamtorch.nn.gpt import GPT, GPTConfig

from .hooks_apps import inspect_forward_pass


def test_forward():
    gpt_config = GPTConfig()
    gpt_model = GPT(gpt_config).cuda()
    example_input = torch.randint(0, gpt_config.vocab_size, (2, 128)).cuda()

    _, layer_spec = inspect_forward_pass(gpt_model, example_input)
    print(layer_spec)


if __name__ == "__main__":
    pytest.main([__file__])
