had the same problem, found the solution here: CUDA out of memory when using Trainer with compute_metrics - #13 by morenolq