tensorrt 1.0.1 → 1.0.2
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- checksums.yaml +4 -4
- data/ext/tensorrt_rb/tensorrt_rb.cpp +15 -1
- data/lib/tensorrt/version.rb +1 -1
- metadata +1 -1
checksums.yaml
CHANGED
|
@@ -1,7 +1,7 @@
|
|
|
1
1
|
---
|
|
2
2
|
SHA256:
|
|
3
|
-
metadata.gz:
|
|
4
|
-
data.tar.gz:
|
|
3
|
+
metadata.gz: 8a0f3db3e3097f404638a13f66c9ce734f2eece4fad78c8cb3a7c7285b9bf26d
|
|
4
|
+
data.tar.gz: e154c61c9e7c38b89851fb4351551232da0c06e6ffba4cb051904c0fd4ade987
|
|
5
5
|
SHA512:
|
|
6
|
-
metadata.gz:
|
|
7
|
-
data.tar.gz:
|
|
6
|
+
metadata.gz: 6105c9e148a53c09426043ab400d3750b2a62e1074512451302447e3a33ffbec8b23c233d8ec93d2f4a8a4ea29a1f43cacfc34a3117f841ac27fe635ca7037ba
|
|
7
|
+
data.tar.gz: a84a24ba7fbe0d2da9ff111d6245b0f46ca7bad7495f6dd2ab364ea1af5c364f07930d54bc440da54f3358f26086a13d73e20b1b528ce552724febb2d4f60e53
|
|
@@ -115,7 +115,21 @@ public:
|
|
|
115
115
|
|
|
116
116
|
size_t get_tensor_bytes(const std::string& name) const {
|
|
117
117
|
auto dims = engine->getTensorShape(name.c_str());
|
|
118
|
-
|
|
118
|
+
auto dtype = engine->getTensorDataType(name.c_str());
|
|
119
|
+
size_t element_size;
|
|
120
|
+
switch (dtype) {
|
|
121
|
+
case nvinfer1::DataType::kINT64: element_size = 8; break;
|
|
122
|
+
case nvinfer1::DataType::kFLOAT:
|
|
123
|
+
case nvinfer1::DataType::kINT32: element_size = 4; break;
|
|
124
|
+
case nvinfer1::DataType::kHALF:
|
|
125
|
+
case nvinfer1::DataType::kBF16: element_size = 2; break;
|
|
126
|
+
case nvinfer1::DataType::kINT8:
|
|
127
|
+
case nvinfer1::DataType::kBOOL:
|
|
128
|
+
case nvinfer1::DataType::kFP8:
|
|
129
|
+
case nvinfer1::DataType::kUINT8: element_size = 1; break;
|
|
130
|
+
default: element_size = 4; break;
|
|
131
|
+
}
|
|
132
|
+
size_t bytes = element_size;
|
|
119
133
|
for (int i = 0; i < dims.nbDims; i++) {
|
|
120
134
|
bytes *= dims.d[i];
|
|
121
135
|
}
|
data/lib/tensorrt/version.rb
CHANGED