How to convert tensor to image array?
I would like to convert a tensor to image array and use tensor.data() method. 

But it doesn't work.

#include &lt;torch/script.h&gt; // One-stop header.

#include &lt;iostream&gt;
#include &lt;memory&gt;
#include &lt;sstream&gt;
#include &lt;string&gt;
#include &lt;vector&gt;

#include "itkImage.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkImageRegionIterator.h"

//////////////////////////////////////////////////////
//Goal: load jit script model and segment myocardium
//Step: 1. load jit script model
//      2. load input image
//      3. predict by model
//      4. save the result to file
//////////////////////////////////////////////////////
typedef short                                               PixelType;
const unsigned int Dimension = 3;
typedef itk::Image&lt;PixelType, Dimension&gt;                    ImageType;
typedef itk::ImageFileReader&lt;ImageType&gt;                     ReaderType;
typedef itk::ImageRegionIterator&lt;ImageType&gt;                 IteratorType;

bool itk2tensor(ImageType::Pointer itk_img, torch::Tensor &amp;tensor_img) {

    typename ImageType::RegionType region = itk_img-&gt;GetLargestPossibleRegion();
    const typename ImageType::SizeType size = region.GetSize();
    std::cout &lt;&lt; "Input size: " &lt;&lt; size[0] &lt;&lt; ", " &lt;&lt; size[1]&lt;&lt; ", " &lt;&lt; size[2] &lt;&lt; std::endl;

    int len = size[0] * size[1] * size[2];
    short rowdata[len];
    int count = 0;
    IteratorType iter(itk_img, itk_img-&gt;GetRequestedRegion());

    // convert itk to array
    for (iter.GoToBegin(); !iter.IsAtEnd(); ++iter) {
        rowdata[count] = iter.Get();
        count++;
    }
    std::cout &lt;&lt; "Convert itk to array DONE!" &lt;&lt; std::endl;

    // convert array to tensor
    tensor_img = torch::from_blob(rowdata, {1, 1, (int)size[0], (int)size[1], (int)size[2]}, torch::kShort).clone();
    tensor_img = tensor_img.toType(torch::kFloat);
    tensor_img = tensor_img.to(torch::kCUDA);
    tensor_img.set_requires_grad(0);

    return true;
}


bool tensor2itk(torch::Tensor &amp;t, ImageType::Pointer itk_img) {

    std::cout &lt;&lt; "tensor dtype = " &lt;&lt; t.dtype() &lt;&lt; std::endl;
    std::cout &lt;&lt; "tensor size = " &lt;&lt; t.sizes() &lt;&lt; std::endl;
    t = t.toType(torch::kShort);
    short * array = t.data&lt;short&gt;();

    ImageType::IndexType start;
    start[0] = 0;  // first index on X
    start[1] = 0;  // first index on Y
    start[2] = 0;  // first index on Z

    ImageType::SizeType  size;
    size[0] = t.size(2);
    size[1] = t.size(3);
    size[2] = t.size(4);

    ImageType::RegionType region;
    region.SetSize( size );
    region.SetIndex( start );

    itk_img-&gt;SetRegions( region );
    itk_img-&gt;Allocate();

    int len = size[0] * size[1] * size[2];

    IteratorType iter(itk_img, itk_img-&gt;GetRequestedRegion());
    int count = 0;
    // convert array to itk
    std::cout &lt;&lt; "start!" &lt;&lt; std::endl;
    for (iter.GoToBegin(); !iter.IsAtEnd(); ++iter) {
        short temp = *array++;    //  ERROR!
        std::cout &lt;&lt; temp &lt;&lt; " ";
        iter.Set(temp);
        count++;
    }
    std::cout &lt;&lt; "end!" &lt;&lt; std::endl;

    return true;
}


int main(int argc, const char* argv[]) {
    int a, b, c;
    if (argc != 4) {
        std::cerr &lt;&lt; "usage: automyo input jitmodel output\n";
        return -1;
    }

    std::cout &lt;&lt; "=========  jit start  =========\n";
    // 1. load jit script model
    std::cout &lt;&lt; "Load script module: " &lt;&lt; argv[2] &lt;&lt; std::endl;
    std::shared_ptr&lt;torch::jit::script::Module&gt; module = torch::jit::load(argv[2]);
    module-&gt;to(at::kCUDA);

    // assert(module != nullptr);
    std::cout &lt;&lt; "Load script module DONE" &lt;&lt; std::endl;

    // 2. load input image
    const char* img_path = argv[1];
    std::cout &lt;&lt; "Load image: " &lt;&lt; img_path &lt;&lt; std::endl;

    ReaderType::Pointer reader = ReaderType::New();

    if (!img_path) {
        std::cout &lt;&lt; "Load input file error!" &lt;&lt; std::endl;
        return false;
    }

    reader-&gt;SetFileName(img_path);
    reader-&gt;Update();

    std::cout &lt;&lt; "Load image DONE!" &lt;&lt; std::endl;

    ImageType::Pointer itk_img = reader-&gt;GetOutput();

    torch::Tensor tensor_img;
    if (!itk2tensor(itk_img, tensor_img)) {
        std::cerr &lt;&lt; "itk2tensor ERROR!" &lt;&lt; std::endl;
    }
    else {
        std::cout &lt;&lt; "Convert array to tensor DONE!" &lt;&lt; std::endl;
    }

    std::vector&lt;torch::jit::IValue&gt; inputs;
    inputs.push_back(tensor_img);

    // 3. predict by model
    torch::Tensor y = module-&gt;forward(inputs).toTensor();
    std::cout &lt;&lt; "Inference DONE!" &lt;&lt; std::endl;

    // 4. save the result to file
    torch::Tensor seg = y.gt(0.5);
    // std::cout &lt;&lt; seg &lt;&lt; std::endl;

    ImageType::Pointer out_itk_img = ImageType::New();
    if (!tensor2itk(seg, out_itk_img)) {
        std::cerr &lt;&lt; "tensor2itk ERROR!" &lt;&lt; std::endl;
    }
    else {
        std::cout &lt;&lt; "Convert tensor to itk DONE!" &lt;&lt; std::endl;
    }

    std::cout &lt;&lt; out_itk_img &lt;&lt; std::endl;

    return true;
}


The runtime log is showed below:


  Load script module:model_myo_jit.pt 
  
  Load script module DONE Load
  
  image: patch_6.nii.gz 
  
  Load image DONE! 
  
  Input size: 128, 128, 128
  
  Convert itk to array DONE! 
  
  Convert array to tensor DONE! 
  
  Inference DONE! 
  
  tensor dtype = unsigned char
  
  tensor size = [1, 1, 96, 96, 96]
  
  start! 
  
  Segmentation fault (core dumped)


Why and how to convert?

I have found the solution. When I convert the y to kCPU, it works. Because it in CUDA before.
