Allow use_decaying_avg_baseline for TraceGraph_ELBO to be used on GPU
@@ -137,7 +137,7 @@ def _compute_elbo_non_reparam(guide_trace, non_reparam_nodes, downstream_costs):
param_name = "__baseline_avg_downstream_cost_" + node
with torch.no_grad():
avg_downstream_cost_old = pyro.param(param_name,
- torch.tensor(0.0).expand(dc_shape).clone())
+ guide_site['value'].new_zeros(dc_shape))
avg_downstream_cost_new = (1 - baseline_beta) * downstream_cost + \
baseline_beta * avg_downstream_cost_old
pyro.get_param_store().replace_param(param_name, avg_downstream_cost_new,
