26 lines
596 B
Python
26 lines
596 B
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import unittest
|
||
|
|
|
||
|
|
from airfrans_frontier.runtime import remove_pythonpath_entries
|
||
|
|
|
||
|
|
remove_pythonpath_entries()
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from airfrans_frontier.models import PointwiseMLP
|
||
|
|
|
||
|
|
|
||
|
|
class PointwiseMLPTests(unittest.TestCase):
|
||
|
|
def test_forward_pass_returns_batch_by_target_dim(self) -> None:
|
||
|
|
model = PointwiseMLP(input_dim=5, output_dim=4, hidden_width=16, depth=2, activation="gelu")
|
||
|
|
batch = torch.randn(7, 5)
|
||
|
|
|
||
|
|
output = model(batch)
|
||
|
|
|
||
|
|
self.assertEqual(tuple(output.shape), (7, 4))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|