@@ -55,6 +55,23 @@ inline void FillSequentialTensor(const std::shared_ptr<Tensor> &tensor, float st
5555 for (size_t i = 0 ; i < size; ++i) { data[i] = start + static_cast <float >(i); }
5656}
5757
58+ inline void FillConstantTensor (const std::shared_ptr<Tensor> &tensor, float value) {
59+ size_t size = 1 ;
60+ for (auto dim : tensor->Dims ()) { size *= static_cast <size_t >(dim); }
61+
62+ if (!tensor->GetDevice ().IsCPU ()) {
63+ auto cpu_tensor
64+ = std::make_shared<Tensor>(tensor->Dims (), tensor->Dtype (), Device (Device::DeviceType::kCPU , 0 ));
65+ auto *cpu_data = static_cast <float *>(cpu_tensor->DataPtr ());
66+ for (size_t i = 0 ; i < size; ++i) { cpu_data[i] = value; }
67+ tensor->CopyFrom (cpu_tensor);
68+ return ;
69+ }
70+
71+ auto *data = static_cast <float *>(tensor->DataPtr ());
72+ for (size_t i = 0 ; i < size; ++i) { data[i] = value; }
73+ }
74+
5875#define REQUIRE_MIN_DEVICES (n ) \
5976 do { \
6077 int available_gpus = infini_train::test::GetCudaDeviceCount (); \
@@ -107,11 +124,14 @@ class AutogradTestBase : public InfiniTrainTest {
107124};
108125
109126inline std::vector<Device::DeviceType> CudaDeviceTypes () {
110- if (HasCudaRuntime ()) {
111- return {Device::DeviceType::kCUDA };
112- }
113- LOG (INFO ) << " No CUDA runtime found, skipping CUDA tests." ;
114- return {};
127+ static const std::vector<Device::DeviceType> types = []() {
128+ if (HasCudaRuntime ()) {
129+ return std::vector<Device::DeviceType>{Device::DeviceType::kCUDA };
130+ }
131+ LOG (INFO ) << " No CUDA runtime found, skipping CUDA tests." ;
132+ return std::vector<Device::DeviceType>{};
133+ }();
134+ return types;
115135}
116136
117137} // namespace test
0 commit comments