The short answer is yes, it's possible and easy. The longer answer is, that's a bit of an apples to oranges comparison. Pytorch is calling optimized c++ kernels

The short answer is yes, it's possible and easy. The longer answer is, that's a bit of an apples to oranges comparison. Pytorch is calling optimized c++ kernels, whereas Flux is completely generic and can do GPU and CPU codegen (including AMD). Flux can do things Pytorch just doesn't have CUDA code for, It's more generic and composable and it's getting there with a much smaller team/budget. That means it has more potential, with the drawback that it isn't quite as fast, for now, in some of the common use cases the pytorch has been hard coded/ hyperoptimized for. In other more complex cases beyond big fat matmuls and such, it's much faster, even comparing to torchscript (sciml, neural ODEs etc). Still, the codegen is getting there with the funded work on the GPU compiler and memory allocation and there eventually won't be any sort of tradeoff. In the meantime, Yes, there is a package at allows calling pytorch kernels without modifying your code. This is also only possible with such finesse due to Julia's power and composability. See here for the package: https://fluxml.ai/2020/06/29/acclerating-flux-torch.html https://fluxml.ai/2020/06/29/acclerating-flux-torch.html