Skip to content

Commit 595b9f4

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,2), instead of [0,5) - Used random float inputs in the range [0, 2) instead of int, then casted to desired type. The ultimate effect of this change with regard to bug #1418, is random floats are selected in the range [0, 2), then casted to Int, effectively making the range of allowed ints {0, 1}, as required by the model - More robust fix to follow
1 parent 7b37ada commit 595b9f4

File tree

1 file changed

+11
-3
lines changed

1 file changed

+11
-3
lines changed

core/partitioning/shape_analysis.cpp

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,16 +12,24 @@ namespace partitioning {
1212
at::Tensor generateSingleInput(ir::Input& input, c10::optional<at::ScalarType>& type_opt) {
1313
auto cur_shape = input.input_shape;
1414
std::vector<int64_t> shape;
15+
16+
// Initialize min and max ranges for random number selection
17+
int LoValIncl = 0;
18+
int HiValExcl = 2;
19+
1520
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];
21+
1722
auto type = at::kFloat;
1823
if (type_opt) {
1924
type = type_opt.value();
2025
} else {
2126
LOG_WARNING("Input type for doing shape analysis could not be determined, defaulting to F32");
2227
}
23-
auto in = at::randint(5, shape, {at::kCUDA}).to(type);
24-
// ivalue_map[input.first] = in.clone();
28+
29+
// Make the value range for input tensor a uniform (float) distribution
30+
// over [LoValIncl, HiValExcl), then cast to the desired dtype
31+
auto in = ((HiValExcl - LoValIncl) * at::rand(shape, {at::kCUDA}) + LoValIncl).to(type);
32+
2533
return in;
2634
}
2735

0 commit comments

Comments
 (0)