It’s done. It seems I still have some problems though, and they might also be related to cuda.
I can’t reproduce the error with trivial tensors, but basically the log_softmax function produces nans and -inf.
I’m not even sure those are “reasonable” values to obtain for a last layer, but in any case the cuda function doesn’t seem happy with them.
last_layer
Variable containing:
(0 ,0 ,.,.) = 
    3.5879    3.6678    3.8380  ...     3.1548    3.0576    2.9584
    3.4753    3.7363    3.8944  ...     2.9736    3.0051    2.9889
    3.3298    3.3160    3.5382  ...     2.8276    2.8111    2.8584
              ...                ⋱                ...             
    3.2416    3.1960    3.2502  ...    90.0304   98.9006   98.5473
    3.2719    3.2843    3.2724  ...    39.1980   67.7482   73.4172
    3.2535    3.3880    3.3061  ...    13.5371   25.7164   37.9838

(0 ,1 ,.,.) = 
   -3.7768   -3.8683   -4.0629  ...    -3.2815   -3.1703   -3.0568
   -3.6481   -3.9466   -4.1274  ...    -3.0742   -3.1102   -3.0917
   -3.4816   -3.4659   -3.7200  ...    -2.9072   -2.8883   -2.9424
              ...                ⋱                ...             
   -3.3808   -3.3285   -3.3906  ...  -122.8841 -137.5700 -122.2672
   -3.4154   -3.4296   -3.4160  ...   -48.6680  -82.5093  -91.6747
   -3.3943   -3.5482   -3.4545  ...   -17.6337  -33.1317  -46.8320
[torch.cuda.FloatTensor of size 1x2x68x68 (GPU 0)]

F.log_softmax(last_layer)
Variable containing:
(0 ,0 ,.,.) = 
   -0.0006   -0.0005   -0.0004  ...    -0.0016   -0.0020   -0.0024
   -0.0008   -0.0005   -0.0003  ...    -0.0024   -0.0022   -0.0023
   -0.0011   -0.0011   -0.0007  ...    -0.0032   -0.0033   -0.0030
              ...                ⋱                ...             
   -0.0013   -0.0015   -0.0013  ...        nan       nan       nan
   -0.0012   -0.0012   -0.0012  ...     0.0000    0.0000    0.0000
   -0.0013   -0.0010   -0.0012  ...     0.0000    0.0000    0.0000

(0 ,1 ,.,.) = 
   -7.3653   -7.5366   -7.9013  ...    -6.4379   -6.2299   -6.0176
   -7.1242   -7.6833   -8.0221  ...    -6.0501   -6.1175   -6.0829
   -6.8125   -6.7830   -7.2590  ...    -5.7381   -5.7027   -5.8038
              ...                ⋱                ...             
   -6.6237   -6.5260   -6.6420  ...       -inf      -inf      -inf
   -6.6886   -6.7152   -6.6896  ...   -87.8659      -inf      -inf
   -6.6491   -6.9371   -6.7618  ...   -31.1708  -58.8481  -84.8159
[torch.cuda.FloatTensor of size 1x2x68x68 (GPU 0)]

F.log_softmax(last_layer.cpu())
Variable containing:
(0 ,0 ,.,.) = 
   -0.0006   -0.0005   -0.0004  ...    -0.0016   -0.0020   -0.0024
   -0.0008   -0.0005   -0.0003  ...    -0.0024   -0.0022   -0.0023
   -0.0011   -0.0011   -0.0007  ...    -0.0032   -0.0033   -0.0030
              ...                ⋱                ...             
   -0.0013   -0.0015   -0.0013  ...     0.0000    0.0000    0.0000
   -0.0012   -0.0012   -0.0012  ...     0.0000    0.0000    0.0000
   -0.0013   -0.0010   -0.0012  ...    -0.0000    0.0000    0.0000

(0 ,1 ,.,.) = 
   -7.3653   -7.5366   -7.9013  ...    -6.4379   -6.2299   -6.0176
   -7.1242   -7.6833   -8.0221  ...    -6.0501   -6.1175   -6.0829
   -6.8125   -6.7830   -7.2590  ...    -5.7381   -5.7027   -5.8038
              ...                ⋱                ...             
   -6.6237   -6.5260   -6.6420  ...  -212.9145 -236.4706 -220.8145
   -6.6886   -6.7152   -6.6896  ...   -87.8659 -150.2575 -165.0919
   -6.6491   -6.9371   -6.7618  ...   -31.1708  -58.8481  -84.8159
[torch.FloatTensor of size 1x2x68x68]