mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
e76fa9822d
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/32607 As desc. Test Plan: Unit-test. Reviewed By: xw285cornell, chocjy Differential Revision: D19551567 fbshipit-source-id: 3a121351d2b4016e99a1536dec746be970698664
41 lines
1.5 KiB
Python
41 lines
1.5 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
from caffe2.python import core, utils, test_util
|
|
|
|
import numpy as np
|
|
|
|
|
|
class TestUtils(test_util.TestCase):
|
|
def testArgsToDict(self):
|
|
args = [utils.MakeArgument("int1", 3),
|
|
utils.MakeArgument("float1", 4.0),
|
|
utils.MakeArgument("string1", "foo"),
|
|
utils.MakeArgument("intlist1", np.array([3, 4])),
|
|
utils.MakeArgument("floatlist1", np.array([5.0, 6.0])),
|
|
utils.MakeArgument("stringlist1", np.array(["foo", "bar"]))]
|
|
dict_ = utils.ArgsToDict(args)
|
|
expected = {"int1" : 3,
|
|
"float1" : 4.0,
|
|
"string1" : b"foo",
|
|
"intlist1" : [3, 4],
|
|
"floatlist1" : [5.0, 6.0],
|
|
"stringlist1" : [b"foo", b"bar"]}
|
|
self.assertEqual(dict_, expected, "dictionary version of arguments "
|
|
"doesn't match original")
|
|
|
|
def testBuildUniqueMutexIter(self):
|
|
init_net = core.Net("init_net")
|
|
net = core.Net("net")
|
|
utils.BuildUniqueMutexIter(init_net, net)
|
|
|
|
for op in init_net.Proto().op:
|
|
self.assertEqual(op.device_option.extra_info[0],
|
|
"device_type_override:cpu")
|
|
|
|
for op in net.Proto().op:
|
|
self.assertEqual(op.device_option.extra_info[0],
|
|
"device_type_override:cpu")
|