torch::data::transforms::Normalize on kCUDA?

C++
1

Hello,
Is it possible to apply a transform on a single image as we do in PyTorch, only directly on the kCUDA tensor? My example below works only if I apply the transform on a tensor that is in the CPU:

        auto input_tensor = torch::from_blob(frame.data, {1, kIMAGE_SIZE, kIMAGE_SIZE, kCHANNELS});
        input_tensor = input_tensor.permute({0, 3, 1, 2});

        std::vector<double> norm_mean = {0.485, 0.456, 0.406};
        std::vector<double> norm_std = {0.229, 0.224, 0.225};
//        auto norm_mean = torch::tensor(
//                {0.485, 0.456, 0.406},
//                torch::requires_grad(false).dtype(torch::kFloat32));
//
//        auto norm_std = torch::tensor(
//                {0.229, 0.224, 0.225},
//                torch::requires_grad(false).dtype(torch::kFloat32));
        input_tensor = torch::data::transforms::Normalize<>(norm_mean, norm_std)(input_tensor);
        input_tensor = input_tensor.to(at::kCUDA);
        torch::Tensor out_tensor = module.forward({input_tensor}).toTensor();

Thanks,

2 Likes