seeker
12 lines
363 B
Python
12 lines
363 B
Python
import torch
|
|
|
|
from alpamayo_lite.diffusion import FlowMatching
|
|
|
|
|
|
def test_flow_matching_target_velocity() -> None:
|
|
flow = FlowMatching((4, 2), inference_steps=2)
|
|
target = torch.ones(1, 4, 2)
|
|
data = flow.construct_training_data(target)
|
|
exact_velocity = data["target_action"] - data["noise"]
|
|
assert float(flow.loss(data, exact_velocity)) == 0.0
|