[Minor] Dataloader jit speedup types

[ourownstory]

Experiments to speed up main Dataloader jit PR

[github-actions]

Model Benchmark

Benchmark	Metric	main	current	diff
YosemiteTemps	MAE_val	0.56412	0.60324	6.94%	⚠️
YosemiteTemps	RMSE_val	0.83161	0.88159	6.01%	⚠️
YosemiteTemps	Loss_val	0.0004	0.00045	12.29%	❌
YosemiteTemps	MAE	0.98449	0.98956	0.52%	✅
YosemiteTemps	RMSE	1.75389	1.73958	-0.82%	✅
YosemiteTemps	Loss	0.00132	0.0013	-1.5%	✅
YosemiteTemps	time	33.9815	841.64	2376.76%	❌
EnergyPriceDaily	MAE_val	5.41157	5.65526	4.5%	⚠️
EnergyPriceDaily	RMSE_val	6.71538	6.86834	2.28%	✅
EnergyPriceDaily	Loss_val	0.02525	0.02645	4.75%	⚠️
EnergyPriceDaily	MAE	5.94936	6.2678	5.35%	⚠️
EnergyPriceDaily	RMSE	7.9833	8.34813	4.57%	⚠️
EnergyPriceDaily	Loss	0.02579	0.02826	9.59%	❌
EnergyPriceDaily	time	15.154	272.94	1701.1%	❌
PeytonManning	MAE_val	0.34955	0.35129	0.5%	✅
PeytonManning	RMSE_val	0.50049	0.50154	0.21%	✅
PeytonManning	Loss_val	0.01771	0.0178	0.53%	✅
PeytonManning	MAE	0.34653	0.34697	0.13%	✅
PeytonManning	RMSE	0.49312	0.49406	0.19%	✅
PeytonManning	Loss	0.01464	0.01468	0.28%	✅
PeytonManning	time	11.8492	273.56	2208.68%	❌
AirPassengers	MAE_val	30.6285	30.9382	1.01%	✅
AirPassengers	RMSE_val	31.5254	32.1915	2.11%	✅
AirPassengers	Loss_val	0.01277	0.01332	4.27%	⚠️
AirPassengers	MAE	6.16861	6.91875	12.16%	❌
AirPassengers	RMSE	7.85225	8.9841	14.41%	❌
AirPassengers	Loss	0.00065	0.00078	20.92%	❌
AirPassengers	time	7.78978	31.07	298.86%	❌

Model training plots

Model Training

PeytonManning

YosemiteTemps

AirPassengers

EnergyPriceDaily