Skip to content

Commit 4ba166e

Browse files
committed
fix: Resolve cuda 710 error when compiling BERT models
- Issue arising when compiling BERT models with 3+ inputs - Added temporary fix by decreasing the range of allowed values to the random number generator for creating input tensors to 0-1, instead of 0-4 - More robust fix to follow
1 parent 5a7f00e commit 4ba166e

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

core/partitioning/shape_analysis.cpp

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,15 +13,15 @@ at::Tensor generateSingleInput(ir::Input& input, c10::optional<at::ScalarType>&
1313
auto cur_shape = input.input_shape;
1414
std::vector<int64_t> shape;
1515
shape.insert(shape.begin(), std::begin(cur_shape.d), std::begin(cur_shape.d) + cur_shape.nbDims);
16-
// auto type_opt = types[input.first][i];
16+
1717
auto type = at::kFloat;
1818
if (type_opt) {
1919
type = type_opt.value();
2020
} else {
2121
LOG_WARNING("Input type for doing shape analysis could not be determined, defaulting to F32");
2222
}
23-
auto in = at::randint(5, shape, {at::kCUDA}).to(type);
24-
// ivalue_map[input.first] = in.clone();
23+
auto in = at::randint(2, shape, {at::kCUDA}).to(type);
24+
2525
return in;
2626
}
2727

0 commit comments

Comments
 (0)