mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
fef52cc1f8
Summary:
This lets you compile builtin functions from C++ without having a dependence on Python
```cpp
auto module = torch::jit::compile(JIT"(
def my_script_method(x, y):
return torch.relu(x) + y
)");
IValue result = module->run_method("my_script_method", 1, 2);
```
goldsborough zdevito apaszke
Pull Request resolved: https://github.com/pytorch/pytorch/pull/10847
Differential Revision: D9543461
Pulled By: driazati
fbshipit-source-id: 6160dae094030ca144a0df93cb9f26aa78c8cf27
32 lines
732 B
C++
32 lines
732 B
C++
#include <catch.hpp>
|
|
|
|
#include <torch/jit.h>
|
|
#include <torch/tensor.h>
|
|
|
|
#include <string>
|
|
|
|
TEST_CASE("torch script") {
|
|
SECTION("multiple functions") {
|
|
auto module = torch::jit::compile(R"JIT(
|
|
def test_mul(a, b):
|
|
return a * b
|
|
def test_relu(a, b):
|
|
return torch.relu(a + b)
|
|
def test_while(a, i):
|
|
while i < 10:
|
|
a += a
|
|
i += 1
|
|
return a
|
|
)JIT");
|
|
auto a = torch::ones(1);
|
|
auto b = torch::ones(1);
|
|
|
|
REQUIRE(1 == module->run_method("test_mul", a, b).toTensor().toCLong());
|
|
|
|
REQUIRE(2 == module->run_method("test_relu", a, b).toTensor().toCLong());
|
|
|
|
REQUIRE(
|
|
0x200 == module->run_method("test_while", a, b).toTensor().toCLong());
|
|
}
|
|
}
|