Seems to be working. I have seen very few examples serving data from memory could you enlighten me as to why? Given that my dataset fits into my ram memory. And even AWS GPU server have ludicrous amounts of ram?
With the changes you suggest I get this
(def calc_loss uses .cpu heavily, any suggestions on how to improve this?)
  2350647 function calls (2293645 primitive calls) in 192.886 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     1050   91.521    0.087   91.521    0.087 {method 'cpu' of 'torch._C._TensorBase' objects}
    11550   16.191    0.001   16.191    0.001 {built-in method conv2d}
      350   10.167    0.029   10.167    0.029 {built-in method binary_cross_entropy_with_logits}
     1750    9.716    0.006    9.716    0.006 {built-in method cat}
    10150    8.492    0.001    8.492    0.001 {built-in method threshold_}
    39682    6.563    0.000    6.563    0.000 {method 'decode' of 'ImagingDecoder' objects}
      700    6.099    0.009    6.099    0.009 {method 'pin_memory' of 'torch._C._TensorBase' objects}
      345    5.686    0.016    5.686    0.016 {method 'run_backward' of 'torch._C._EngineBase' objects}
      350    3.926    0.011    5.144    0.015 Loss.py:5(dice_loss)
     1750    3.245    0.002    3.245    0.002 {built-in method torch._C._nn.upsample_bilinear2d}
     3135    2.887    0.001    2.887    0.001 {method 'convert' of 'ImagingCore' objects}
      700    2.474    0.004    2.474    0.004 {built-in method stack}
    59340    2.422    0.000    2.422    0.000 {method 'mul_' of 'torch._C._TensorBase' objects}
    59340    2.287    0.000    2.287    0.000 {method 'add_' of 'torch._C._TensorBase' objects}
      350    1.842    0.005  108.726    0.311 Trainer.py:79(calc_loss)
     2800    1.772    0.001    1.772    0.001 {method 'resize' of 'ImagingCore' objects}
      968    1.755    0.002    1.845    0.002 {method 'to' of 'torch._C._TensorBase' objects}
    29670    1.700    0.000    1.700    0.000 {method 'sqrt' of 'torch._C._TensorBase' objects}
     2100    1.149    0.001    1.149    0.001 {method 'sum' of 'torch._C._TensorBase' objects}
     7000    1.102    0.000    1.102    0.000 {built-in method batch_norm}
    29670    1.074    0.000    1.074    0.000 {method 'addcdiv_' of 'torch._C._TensorBase' objects}
    29670    1.041    0.000    1.041    0.000 {method 'addcmul_' of 'torch._C._TensorBase' objects}

For reference without
     2303784 function calls (2249456 primitive calls) in 193.183 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
      968  108.155    0.112  108.257    0.112 {method 'to' of 'torch._C._TensorBase' objects}
     1050   42.901    0.041   42.901    0.041 {method 'cpu' of 'torch._C._TensorBase' objects}
    39682    6.515    0.000    6.515    0.000 {method 'decode' of 'ImagingDecoder' objects}
      345    5.392    0.016    5.392    0.016 {method 'run_backward' of 'torch._C._EngineBase' objects}
    11550    2.967    0.000    2.967    0.000 {built-in method conv2d}
     3135    2.848    0.001    2.848    0.001 {method 'convert' of 'ImagingCore' objects}
      700    2.494    0.004    2.494    0.004 {built-in method stack}
    59340    2.389    0.000    2.389    0.000 {method 'mul_' of 'torch._C._TensorBase' objects}
    59340    2.281    0.000    2.281    0.000 {method 'add_' of 'torch._C._TensorBase' objects}
     2800    1.753    0.001    1.753    0.001 {method 'resize' of 'ImagingCore' objects}
    29670    1.704    0.000    1.704    0.000 {method 'sqrt' of 'torch._C._TensorBase' objects}
    29670    1.079    0.000    1.079    0.000 {method 'addcdiv_' of 'torch._C._TensorBase' objects}
     7000    1.044    0.000    1.044    0.000 {built-in method batch_norm}
    29670    1.041    0.000    1.041    0.000 {method 'addcmul_' of 'torch._C._TensorBase' objects}