我正在开发一种算法,它涉及一个CPU情况,允许NumPy,以及一个GPU情况,允许PyTorch。物体几乎总是四维的。该对象的两个版本如下所示。array([0.5625, 1.250, 0.6250])
max_tensor_fn(B_tensor) # returns tensor([0.5625, 1.250, 0.6250]) 这里讨论了部分解决方案,但这仅适用于CPU上的NumPy: Max of each 2D matrix in 4D NumPy array 然而,在图形处理器上,P