@ptrblck, the .csv file can be read from GitHub, but there are still some modules and the net state_dict that need to be downloaded to run the below reduced code version that still demonstrates the CPU and GPU discrepancy issue. If you don’t have time for the additional downloading, I’ll completely understand. The reduced code is as follows:
#Python-related imports
import math
import sys
from datetime import datetime
import os.path

#Torch imports
import torch
import torch.distributions as D
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Function

#PyData imports
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt

#Module imports
from obs_and_flow import *
from plotting import *

#PyTorch settings
if torch.cuda.is_available():
    print('CUDA device detected.')
    active_device = torch.device('cuda')
    torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
    active_device = torch.device('cpu')
    torch.set_default_tensor_type('torch.FloatTensor')    

print(active_device)

torch.set_printoptions(precision = 8)
torch.manual_seed(0)
#IAF SSM time parameters
dt_flow = 1.0 #Increased from 0.1 to reduce memory.
t = 5000 #In hours.
n = int(t / dt_flow) + 1
t_span = np.linspace(0, t, n)
t_span_tensor = torch.reshape(torch.Tensor(t_span), [1, n, 1]).to(active_device) #T_span needs to be converted to tensor object. Additionally, facilitates conversion of I_S and I_D to tensor objects.

#Training parameters
elbo_iter = 105000
elbo_lr = 1e-2
elbo_lr_decay = 0.7
elbo_lr_decay_step_size = 5000
elbo_warmup_iter = 5000
elbo_warmup_lr = 1e-6
ptrain_iter = 0
ptrain_alg = 'L1'
batch_size = 31
eval_batch_size = 31
obs_error_scale = 0.1
prior_scale_factor = 0.25
num_layers = 5
reverse = False
base_state = False

#Specify desired SBM SDE model type and details.
state_dim_SCON = 3
SBM_SDE_class = 'SCON'
diffusion_type = 'C'
learn_CO2 = False
fix_theta_dict = None

now_string = 'SCON-C_no_CO2_logit_short_2022_01_30_16_14_49_discuss'
outputs_folder = 'training_pt_outputs/'
plots_folder = 'training_plots/'
save_string = '_iter_105000_warmup_5000_t_5000_dt_1.0_batch_31_layers_5_lr_0.01_decay_step_5000_warmup_lr_1e-06_sd_scale_0.25_SCON-C_no_CO2_logit_short_2022_01_30_16_14_49.pt'

net_state_dict_save_string = os.path.join(outputs_folder, 'net_state_dict' + save_string)

obs_times, obs_means, obs_error = csv_to_obs_df('https://raw.githubusercontent.com/wallyxie/varInferenceSoilBiogeoModelSyntheticData/main/pytorch_sbm_sde_vi/generated_data/SCON-C_no_CO2_logit_short_2022_01_20_08_53_sample_y_t_5000_dt_0-01_sd_scale_0-25.csv', state_dim_SCON, t, obs_error_scale)
obs_model = ObsModel(DEVICE = active_device, TIMES = obs_times, DT = dt_flow, MU = obs_means, SCALE = obs_error)
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))

#Plot training posterior results and ELBO history.
net.eval()
x, _ = net(eval_batch_size)

plot_states_post(x, None, obs_model, None, elbo_iter, elbo_warmup_iter, t, dt_flow, batch_size, eval_batch_size, num_layers, elbo_lr, elbo_lr_decay_step_size, elbo_warmup_lr, prior_scale_factor, plots_folder, now_string, fix_theta_dict, learn_CO2, ymin_list = [0, 0, 0, 0], ymax_list = [70., 8., 11., 0.025])

At this point, the only thing being loaded is the model state_dict after the model is newly instantiated. The state_dict in question can be downloaded at the following link: varInferenceSoilBiogeoModelSyntheticData/net_iter_105000_warmup_5000_t_5000_dt_1.0_batch_31_layers_5_lr_0.01_decay_step_5000_warmup_lr_1e-06_sd_scale_0.25_SCON-C_no_CO2_logit_short_2022_01_30_16_14_49.pt at main · wallyxie/varInferenceSoilBiogeoModelSyntheticData · GitHub
The custom modules that need to be downloaded to the same directory as the script are “obs_and_flow.py” and “plotting.py.”
When launched on a GPU, the script produces the following plot:

net_iter_105000_warmup_5000_t_5000_dt_1.0_batch_31_samples_31_layers_5_lr_0.01_decay_step_5000_warmup_lr_1e-06_sd_scale_0.25_SCON-C_no_CO2_logit_short_2022_01_30_16_14_49_discuss1920×1440 135 KB

When launched on a CPU, the script produces the following plot:

net_iter_105000_warmup_5000_t_5000_dt_1.0_batch_31_samples_31_layers_5_lr_0.01_decay_step_5000_warmup_lr_1e-06_sd_scale_0.25_SCON-C_no_CO2_logit_short_2022_01_30_16_14_49_discuss_CPU1920×1440 199 KB

I greatly appreciate your patience, @ptrblck, and again, will completely understand if you don’t have the time to get into downloading of additional custom modules.