gentle-tomato-480
09/03/2025, 10:30 AMtorch.Tensor
that is an input to another task.
I notice that when regardless if the tensor is on the GPU or CPU when it's outputted, it's put back on the GPU in the next task when running locally.
This is a bit unexepected behaviour, as I would expect it deserialize/load in the next step on the device it was serialized on (i.e. if the tensor was on CPU when outputted, I expect it to be loaded as a CPU tensor. If the tensor is on GPU, I expect it to be loaded on GPU tensor).
Running flytekit==1.15.3 (locally)gentle-tomato-480
09/03/2025, 10:32 AM