<filename>torch_glow/tests/nodes/bmm_test.py
from __future__ import absolute_import, division, print_function, unicode_literals

import torch

from tests.utils import jitVsGlow


def test_bmm():
    """Basic test of the PyTorch bmm Node on Glow."""

    def test_f(a, b):
        return (a + a).bmm(b)

    x = torch.randn(6, 4, 10)
    y = torch.randn(6, 10, 2)

    jitVsGlow(test_f, x, y, expected_fused_ops={"aten::bmm"})

<filename>torch_glow/tests/nodes/bmm_test.py
from __future__ import absolute_import, division, print_function, unicode_literals

import torch

from tests.utils import jitVsGlow


def test_bmm():
    """Basic test of the PyTorch bmm Node on Glow."""

    def test_f(a, b):
        return (a + a).bmm(b)

    x = torch.randn(6, 4, 10)
    y = torch.randn(6, 10, 2)

    jitVsGlow(test_f, x, y, expected_fused_ops={"aten::bmm"})

Basic test of the PyTorch bmm Node on Glow.