After updating PyTorch on both my CPU and GPU cluster to 1.10.2 (with GPU PyTorch being installed with cudatoolkit = 10.2 since cudatoolkit = 11.3 does not work on my institution’s cluster), the net.eval discrepancy between CPU and GPU persists, and it seems that the earlier pseudo-improvement of the CPU plots has not been observed in my latest trial.
The CPU plot:

net_iter_105000_warmup_1000_t_5000_dt_1.0_batch_31_samples_31_layers_5_lr_0.01_decay_step_4000_warmup_lr_1e-06_sd_scale_0.25_SCON-C_CO2_logit_short_2022_01_29_15_28_37_CPU_testing1920×1440 224 KB

The GPU plot:

net_iter_105000_warmup_1000_t_5000_dt_1.0_batch_31_samples_31_layers_5_lr_0.01_decay_step_4000_warmup_lr_1e-06_sd_scale_0.25_SCON-C_CO2_logit_short_2022_01_29_15_28_37_testing1920×1440 157 KB

The model code and the loading steps are the exact same. The only thing different between the CPU and GPU evaluation plotting codes at this point is the map_location for the torch.load(...) calls. The CPU results at least are consistently wrong for me for multiple devices I’ve tested that are now using the same PyTorch version (1.10.2). The shared torch.load leading up to the plotting is along the lines of
obs_model = torch.load(obs_model_save_string, map_location = active_device)
net = SDEFlow(active_device, obs_model, state_dim_SCON, t, dt_flow, n, NUM_LAYERS = num_layers, REVERSE = reverse, BASE_STATE = base_state)
net.load_state_dict(torch.load(net_state_dict_save_string, map_location = active_device))
p_theta = torch.load(p_theta_save_string, map_location = active_device)
q_theta = torch.load(q_theta_save_string, map_location = active_device)
SBM_SDE_instance = torch.load(SBM_SDE_instance_save_string, map_location = active_device)

net.eval()
x, _ = net(eval_batch_size)

plot_func(...)

I made sure I was still saving the state_dict with  _use_new_zipfile_serialization = False.