Using GPU Computation in LearnDistribution

Question

LearnDistribution has an option for RealNVP, which is directly training a neural network for the distribution. It also has an option for TargetDevice, however, it appears to do nothing. Take the following code:

digits08 = Cases[ResourceData["MNIST"], Except[_ -> 0]][[All, 1]];
ldtrial4 = 
 LearnDistribution[digits08, 
  Method -> {"RealNVP", "NetworkDepth" -> 2, 
    "CouplingLayersNumber" -> 2, MaxTrainingRounds -> 100}, 
  PerformanceGoal -> "DirectTraining",TargetDevice->"GPU"]

The progress paclet is reporting to me that it's using CPU, instead of GPU. If I try any NetTrain with TargetDevice GPU it works absolutely fine, but LearnDistribution for some reason is not. Is there a way around this? Is this happening on other people's machines?

Answer 1

It was simpler than I thought - simply run:

SetOptions[NetTrain,TargetDevice->"GPU"]

You can't specify GPU training for LearnDistribution, apparently, but you can set the default to be GPU!