mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-15 21:00:47 +00:00
5eb836880d
Summary:
Needed by oss.
This is done by running the following line:
find . -name "*_test.py" -exec sed -i '$ a \\nif __name__ == "__main__":\n import unittest\n unittest.main()' {} \;
Reviewed By: ajtulloch
Differential Revision: D4223848
fbshipit-source-id: ef4696e9701d45962134841165c53e76a2e19233
24 lines
811 B
Python
24 lines
811 B
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
import numpy as np
|
|
|
|
from caffe2.python import core, workspace
|
|
from caffe2.python.test_util import TestCase
|
|
|
|
|
|
class TestGatherOps(TestCase):
|
|
def test_gather_ops(self):
|
|
data = np.array(["world", "hello", "!"], dtype='|S')
|
|
ind = np.array([1, 0, 2], dtype=np.int32)
|
|
workspace.FeedBlob('data', data)
|
|
workspace.FeedBlob('ind', ind)
|
|
workspace.RunOperatorOnce(core.CreateOperator(
|
|
'Gather', ['data', 'ind'], ['word']))
|
|
outdata = np.array(["hello", "world", "!"], dtype='|S')
|
|
assert((workspace.FetchBlob('word') == outdata).all())
|
|
|
|
if __name__ == "__main__":
|
|
import unittest
|
|
unittest.main()
|