@ptrblck I tried with torch.utils.bottleneck and following is the output:
Note: I ran the script for 1 epoch only, as you suggested here 2, and this is printed without starting the training (i.e., before the problem occurred).
--------------------------------------------------------------------------------
  Environment Summary
--------------------------------------------------------------------------------
PyTorch 1.9.0+cu111 DEBUG compiled w/ CUDA 11.1
Running with Python 3.6 and 

`pip3 list` truncated output:
numpy==1.18.1
torch==1.9.0+cu111
torchaudio==0.9.0
torchvision==0.10.0+cu111
--------------------------------------------------------------------------------
  cProfile output
--------------------------------------------------------------------------------
         15127188 function calls (15111211 primitive calls) in 8.870 seconds

   Ordered by: internal time
   List reduced from 2747 to 15 due to restriction <15>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
      289    3.666    0.013    3.666    0.013 {method 'cuda' of 'torch._C._TensorBase' objects}
  1263775    1.174    0.000    1.905    0.000 /usr/lib/python3.6/posixpath.py:75(join)
        1    1.069    1.069    4.304    4.304 /DatasetLoader.py:114(__init__)
  2442625    0.395    0.000    0.395    0.000 {method 'split' of 'str' objects}
  1264426    0.223    0.000    0.365    0.000 /usr/lib/python3.6/posixpath.py:41(_get_sep)
  2560595    0.171    0.000    0.171    0.000 {method 'append' of 'list' objects}
1304392/1303795    0.149    0.000    0.155    0.000 {built-in method builtins.isinstance}
   173823    0.144    0.000    0.151    0.000 /usr/lib/python3.6/glob.py:114(_iterdir)
  1265133    0.128    0.000    0.128    0.000 {method 'endswith' of 'str' objects}
  1270598    0.127    0.000    0.127    0.000 {method 'startswith' of 'str' objects}
        1    0.117    0.117    0.288    0.288 /DatasetLoader.py:129(<listcomp>)
  1268919    0.113    0.000    0.113    0.000 {built-in method posix.fspath}
        2    0.110    0.055    0.115    0.057 {method 'readlines' of '_io._IOBase' objects}
  1092599    0.091    0.000    0.091    0.000 {method 'strip' of 'str' objects}
        1    0.066    0.066    0.775    0.775 /DatasetLoader.py:62(__init__)


--------------------------------------------------------------------------------
  autograd profiler output (CPU mode)
--------------------------------------------------------------------------------
        top 15 events sorted by cpu_time_total

------------------  ------------  ------------  ------------  ------------  ------------  ------------  
              Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
------------------  ------------  ------------  ------------  ------------  ------------  ------------  
     aten::normal_         9.89%       1.233ms         9.89%       1.233ms       1.233ms             1  
     aten::normal_         9.61%       1.198ms         9.61%       1.198ms       1.198ms             1  
     aten::normal_         9.54%       1.190ms         9.54%       1.190ms       1.190ms             1  
     aten::normal_         9.48%       1.182ms         9.48%       1.182ms       1.182ms             1  
     aten::normal_         9.42%       1.174ms         9.42%       1.174ms       1.174ms             1  
    aten::uniform_         6.81%     848.913us         6.81%     848.913us     848.913us             1  
    aten::uniform_         6.35%     791.197us         6.35%     791.197us     791.197us             1  
    aten::uniform_         6.26%     780.781us         6.26%     780.781us     780.781us             1  
    aten::uniform_         5.96%     743.103us         5.96%     743.103us     743.103us             1  
    aten::uniform_         5.94%     740.487us         5.94%     740.487us     740.487us             1  
    aten::uniform_         5.64%     703.088us         5.64%     703.088us     703.088us             1  
     aten::normal_         5.29%     660.099us         5.29%     660.099us     660.099us             1  
         aten::max         3.30%     411.601us         3.43%     427.419us     427.419us             1  
     aten::normal_         3.29%     410.359us         3.29%     410.359us     410.359us             1  
     aten::normal_         3.21%     400.519us         3.21%     400.519us     400.519us             1  
------------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 12.467ms

--------------------------------------------------------------------------------
  autograd profiler output (CUDA mode)
--------------------------------------------------------------------------------
        top 15 events sorted by cpu_time_total

	Because the autograd profiler uses the CUDA event API,
	the CUDA time column reports approximately max(cuda_time, cpu_time).
	Please ignore this output if your code does not use CUDA.

------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
              Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  
------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
     aten::normal_         9.13%       1.198ms         9.13%       1.198ms       1.198ms       1.198ms         9.13%       1.198ms       1.198ms             1  
     aten::normal_         8.93%       1.172ms         8.93%       1.172ms       1.172ms       1.173ms         8.94%       1.173ms       1.173ms             1  
     aten::normal_         8.93%       1.172ms         8.93%       1.172ms       1.172ms       1.171ms         8.93%       1.171ms       1.171ms             1  
     aten::normal_         8.86%       1.163ms         8.86%       1.163ms       1.163ms       1.163ms         8.86%       1.163ms       1.163ms             1  
     aten::normal_         8.85%       1.161ms         8.85%       1.161ms       1.161ms       1.162ms         8.85%       1.162ms       1.162ms             1  
    aten::uniform_         7.56%     992.395us         7.56%     992.395us     992.395us     991.776us         7.56%     991.776us     991.776us             1  
    aten::uniform_         7.36%     966.001us         7.36%     966.001us     966.001us     966.240us         7.36%     966.240us     966.240us             1  
    aten::uniform_         7.30%     958.179us         7.30%     958.179us     958.179us     957.120us         7.29%     957.120us     957.120us             1  
    aten::uniform_         7.29%     956.692us         7.29%     956.692us     956.692us     957.760us         7.30%     957.760us     957.760us             1  
    aten::uniform_         7.10%     932.259us         7.10%     932.259us     932.259us     929.216us         7.08%     929.216us     929.216us             1  
    aten::uniform_         5.43%     711.949us         5.43%     711.949us     711.949us     712.384us         5.43%     712.384us     712.384us             1  
     aten::normal_         4.34%     569.383us         4.34%     569.383us     569.383us     569.248us         4.34%     569.248us     569.248us             1  
    aten::uniform_         3.49%     457.421us         3.49%     457.421us     457.421us     457.216us         3.48%     457.216us     457.216us             1  
         aten::max         2.48%     324.992us         2.95%     387.751us     387.751us     328.032us         2.50%     387.328us     387.328us             1  
     aten::normal_         2.95%     386.638us         2.95%     386.638us     386.638us     387.072us         2.95%     387.072us     387.072us             1  
------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 13.123ms
Self CUDA time total: 13.123ms