Amortized optimization for computing optimal transport maps