Closed cacerca closed 1 year ago
@cacerca
It looks like the only function that gets pushed to the GPU is _subtract_mean_and_normalize (the call to torch.matmul() explicitly invokes cpu()). How much of a speed-up does this get us?
_subtract_mean_and_normalize
torch.matmul()
cpu()
@cacerca
It looks like the only function that gets pushed to the GPU is
_subtract_mean_and_normalize
(the call totorch.matmul()
explicitly invokescpu()
). How much of a speed-up does this get us?