import logging
import unittest

from .helpers import run_module

logging.basicConfig(level=logging.DEBUG)


class TorchTests(unittest.TestCase):
    def setUp(self):
        self.name = "benchmarker"
        self.imgnet_args = [
            "--framework=pytorch",
            "--problem_size=4",
            "--batch_size=2",
        ]

    def test_vgg16(self):
        run_module(self.name, "--problem=vgg16", *self.imgnet_args)

    def test_resnet50(self):
        run_module(self.name, "--problem=resnet50", *self.imgnet_args)

    def test_conv1d(self):
        # vatai: Strange that only "--problem conv1d",
        # "--problem=conv1d" doesn't.
        run_module(
            self.name,
            "--framework=pytorch",
            "--problem conv1d",
            "--problem_size='(4, 4, 4)'",
            "--batch_size=64",
            "--mode=inference",
        )

    def test_conv2d(self):
        run_module(
            self.name,
            "--framework=pytorch",
            "--problem conv2d",
            "--problem_size='(4, 4, 4, 4)'",
            "--batch_size=64",
            "--mode=inference",
        )

    def test_conv2d_1(self):
        run_module(
            self.name,
            "--framework=pytorch",
            "--problem conv2d_1",
            "--problem_size='(4, 4, 4, 4)'",
            "--batch_size=64",
            "--mode=inference",
        )

    def test_conv2d_2(self):
        run_module(
            self.name,
            "--framework=pytorch",
            "--problem conv2d_2",
            "--problem_size='(4, 4, 4, 4)'",
            "--batch_size=64",
            "--mode=inference",
        )

    def test_conv3d(self):
        run_module(
            self.name,
            "--framework=pytorch",
            "--problem conv3d",
            "--problem_size='(4, 4, 4, 4, 4)'",
            "--batch_size=64",
            "--mode=inference",
        )

    def test_conv3d_1(self):
        run_module(
            self.name,
            "--framework=pytorch",
            "--problem conv3d_1",
            "--problem_size='(4, 4, 4, 4, 4)'",
            "--batch_size=64",
            "--mode=inference",
        )

@@ -816,37 +816,36 @@
   %22--batch_size=
-6
 4%22,%0A
+
             %22--m
@@ -1071,33 +1071,32 @@
   %22--batch_size=
-6
 4%22,%0A            
@@ -1330,33 +1330,32 @@
   %22--batch_size=
-6
 4%22,%0A            
@@ -1589,33 +1589,32 @@
   %22--batch_size=
-6
 4%22,%0A            
@@ -1847,33 +1847,32 @@
   %22--batch_size=
-6
 4%22,%0A            
@@ -2083,32 +2083,32 @@
  4, 4, 4, 4)'%22,%0A
+
             %22--b
@@ -2117,17 +2117,16 @@
 ch_size=
-6
 4%22,%0A    

Fix batch_sizes in test_pytorch.py