ae15f35f07
* Add is_nonzero operator and test * fix * Update is_nonzero.cpp * Update is_nonzero.cpp * requested changes * Update is_nonzero.cpp * Update is_nonzero.cpp --------- Co-authored-by: Maxim Vafin <maxim.vafin@intel.com>
33 lines
967 B
Python
33 lines
967 B
Python
# Copyright (C) 2018-2023 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from pytorch_layer_test_class import PytorchLayerTest
|
|
|
|
|
|
@pytest.mark.parametrize('input_tensor', (np.array([0.]), np.array([1.5]), np.array([False]), np.array([3])))
|
|
class TestIsNonZero(PytorchLayerTest):
|
|
|
|
def _prepare_input(self):
|
|
input_tensor = self.input_tensor
|
|
return (input_tensor.astype(np.int64),)
|
|
|
|
def create_model(self):
|
|
class aten_is_nonzero(torch.nn.Module):
|
|
|
|
def forward(self, input_tensor):
|
|
return torch.is_nonzero(input_tensor)
|
|
|
|
ref_net = None
|
|
|
|
return aten_is_nonzero(), ref_net, "aten::is_nonzero"
|
|
|
|
@pytest.mark.nightly
|
|
@pytest.mark.precommit
|
|
def test_is_nonzero(self, ie_device, precision, ir_version, input_tensor):
|
|
self.input_tensor = input_tensor
|
|
self._test(*self.create_model(), ie_device, precision, ir_version)
|