When using bmm, it can output right result on cpu, but always 0s on device, which is not expected. the code is as follows:
void test_bmm() {
auto t1 = torch::randn({1, 3, 4});
auto t2 = torch::randn({1, 4, 2});
auto r = torch::bmm(t1, t2);
std::cout << r << std::endl; // Right
t1 = t1.cuda();
t2 = t2.cuda();
auto r2 = torch::bmm(t1, t2);
std::cout << r2 << std::endl; // Always 0s
/*
(1,.,.) =
0  0
0  0
0  0
[ CUDAFloatType{1,3,2} ]
*/
}
I want to know what’s wrong with it?
Thanks