🩹 Fix missing import

probabilists · May 24, 2022 · 40f6a94 · 40f6a94
1 parent 53df6a7
commit 40f6a94
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 4 deletions.
diff --git a/lampe/inference.py b/lampe/inference.py
@@ -10,6 +10,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+from itertools import islice
 from torch import Tensor, BoolTensor, Size
 from typing import *
 

diff --git a/notebooks/01_npe.ipynb b/notebooks/01_npe.ipynb
@@ -84,7 +84,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "loader = JointLoader(prior, sim, numpy=False, batched=True)"
+    "loader = JointLoader(prior, sim, batch_size=2**10, numpy=False, batched=True)"
    ]
   },
   {

diff --git a/notebooks/02_nre.ipynb b/notebooks/02_nre.ipynb
@@ -83,7 +83,7 @@
     "    loader = JointLoader(prior, sim, batched=True)\n",
     "    H5Dataset.store(loader, 'train.h5', size=2**20)  # store 1048576 pairs on disk\n",
     "\n",
-    "dataset = H5Dataset('train.h5', batch_size=2**10, pin_memory=True)\n",
+    "dataset = H5Dataset('train.h5', batch_size=2**10, shuffle=True, pin_memory=True)\n",
     "dataset[0]"
    ]
   },
@@ -203,7 +203,7 @@
     "x_star = sim(theta_star).cuda()\n",
     "\n",
     "estimator.eval()\n",
-    "deepto(prior, device='cuda')\n",
+    "deepto(prior, device='cuda')  # move internal tensors to CUDA\n",
     "\n",
     "with torch.no_grad():\n",
     "    theta_0 = prior.sample((2**12,))  # 4096 concurrent Markov chains\n",

diff --git a/setup.py b/setup.py
@@ -10,7 +10,7 @@
 
 setuptools.setup(
     name='lampe',
-    version='0.4.1',
+    version='0.4.2',
     packages=setuptools.find_packages(),
     description='Likelihood-free AMortized Posterior Estimation with PyTorch',
     keywords='parameter inference bayes posterior amortized likelihood ratio mcmc torch',