def _get_test_model_path(name):
sub_dir, file = TEST_MODELS[name]
if sub_dir == "FUSION":
return get_fusion_test_model(file)
else:
return get_test_data_path(sub_dir, file)
def test_megatron_gpt2_attention_fusion(self):
path = get_test_data_path("models", "gpt2_megatron.onnx")
model = onnx.load(path)
optimized_model = optimize_by_fusion(model, model_type="gpt2")
self.verify_fusion(optimized_model, "gpt2_megatron_opt.onnx")