import logging
import unittest

from ..helpers import run_module

logging.basicConfig(level=logging.DEBUG)


class PytorchBertTest(unittest.TestCase):
    def setUp(self):
        self.name = "benchmarker"

    def test_bert(self):
        run_module(
            self.name,
            "--framework=pytorch",
            "--problem=bert",
            "--problem_size=2,8",
            "--batch_size=8",
            "--nb_epoch=1",
            "--mode=inference",
        )

import logging
import unittest

from ..helpers import run_module

logging.basicConfig(level=logging.DEBUG)


class PytorchBertTest(unittest.TestCase):
    def setUp(self):
        self.name = "benchmarker"

    def test_bert(self):
        run_module(
            self.name,
            "--framework=pytorch",
            "--problem=bert",
            "--problem_size=4,8",
            "--batch_size=2",
            "--nb_epoch=1",
            "--mode=inference",
        )

fix batch size