We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Training SNRE_A on GPU fails
Hi! Using SNRE_A with `device="gpu"` currently fails :-(

The error is as follows:
```
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
```
The origin of the issue is in `SNRE_A._loss` which instantiates `labels` without moving it to the device. Adding 
```
labels = labels.to(self._device)
```
below line 126 of `sbi/inference/snre/snre_a.py` fixes the issue.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `sbi/inference/snre/snre_a.py`
Content:
```
1 from typing import Any, Callable, Dict, Optional, Union
2 
3 import torch
4 from torch import Tensor, nn, ones
5 
6 from sbi.inference.posteriors.base_posterior import NeuralPosterior
7 from sbi.inference.snre.snre_base import RatioEstimator
8 from sbi.types import TensorboardSummaryWriter
9 from sbi.utils import del_entries
10 
11 
12 class SNRE_A(RatioEstimator):
13     def __init__(
14         self,
15         prior,
16         classifier: Union[str, Callable] = "resnet",
17         device: str = "cpu",
18         logging_level: Union[int, str] = "warning",
19         summary_writer: Optional[TensorboardSummaryWriter] = None,
20         show_progress_bars: bool = True,
21         **unused_args
22     ):
23         r"""AALR[1], here known as SNRE_A.
24 
25         [1] _Likelihood-free MCMC with Amortized Approximate Likelihood Ratios_, Hermans
26             et al., ICML 2020, https://arxiv.org/abs/1903.04057
27 
28         Args:
29             prior: A probability distribution that expresses prior knowledge about the
30                 parameters, e.g. which ranges are meaningful for them. Any
31                 object with `.log_prob()`and `.sample()` (for example, a PyTorch
32                 distribution) can be used.
33             classifier: Classifier trained to approximate likelihood ratios. If it is
34                 a string, use a pre-configured network of the provided type (one of
35                 linear, mlp, resnet). Alternatively, a function that builds a custom
36                 neural network can be provided. The function will be called with the
37                 first batch of simulations (theta, x), which can thus be used for shape
38                 inference and potentially for z-scoring. It needs to return a PyTorch
39                 `nn.Module` implementing the classifier.
40             device: torch device on which to compute, e.g. gpu, cpu.
41             logging_level: Minimum severity of messages to log. One of the strings
42                 INFO, WARNING, DEBUG, ERROR and CRITICAL.
43             summary_writer: A tensorboard `SummaryWriter` to control, among others, log
44                 file location (default is `<current working directory>/logs`.)
45             show_progress_bars: Whether to show a progressbar during simulation and
46                 sampling.
47             unused_args: Absorbs additional arguments. No entries will be used. If it
48                 is not empty, we warn. In future versions, when the new interface of
49                 0.14.0 is more mature, we will remove this argument.
50         """
51 
52         kwargs = del_entries(locals(), entries=("self", "__class__", "unused_args"))
53         super().__init__(**kwargs, **unused_args)
54 
55     def train(
56         self,
57         training_batch_size: int = 50,
58         learning_rate: float = 5e-4,
59         validation_fraction: float = 0.1,
60         stop_after_epochs: int = 20,
61         max_num_epochs: Optional[int] = None,
62         clip_max_norm: Optional[float] = 5.0,
63         exclude_invalid_x: bool = True,
64         resume_training: bool = False,
65         discard_prior_samples: bool = False,
66         retrain_from_scratch_each_round: bool = False,
67         show_train_summary: bool = False,
68     ) -> NeuralPosterior:
69         r"""
70         Return classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
71 
72         Args:
73             training_batch_size: Training batch size.
74             learning_rate: Learning rate for Adam optimizer.
75             validation_fraction: The fraction of data to use for validation.
76             stop_after_epochs: The number of epochs to wait for improvement on the
77                 validation set before terminating training.
78             max_num_epochs: Maximum number of epochs to run. If reached, we stop
79                 training even when the validation loss is still decreasing. If None, we
80                 train until validation loss increases (see also `stop_after_epochs`).
81             clip_max_norm: Value at which to clip the total gradient norm in order to
82                 prevent exploding gradients. Use None for no clipping.
83             exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=±∞`
84                 during training. Expect errors, silent or explicit, when `False`.
85             resume_training: Can be used in case training time is limited, e.g. on a
86                 cluster. If `True`, the split between train and validation set, the
87                 optimizer, the number of epochs, and the best validation log-prob will
88                 be restored from the last time `.train()` was called.
89             discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
90                 from the prior. Training may be sped up by ignoring such less targeted
91                 samples.
92             retrain_from_scratch_each_round: Whether to retrain the conditional density
93                 estimator for the posterior from scratch each round.
94             show_train_summary: Whether to print the number of epochs and validation
95                 loss and leakage after the training.
96 
97         Returns:
98             Classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
99         """
100 
101         # AALR is defined for `num_atoms=2`.
102         # Proxy to `super().__call__` to ensure right parameter.
103         kwargs = del_entries(locals(), entries=("self", "__class__"))
104         return super().train(**kwargs, num_atoms=2)
105 
106     def _loss(self, theta: Tensor, x: Tensor, num_atoms: int) -> Tensor:
107         """
108         Returns the binary cross-entropy loss for the trained classifier.
109 
110         The classifier takes as input a $(\theta,x)$ pair. It is trained to predict 1
111         if the pair was sampled from the joint $p(\theta,x)$, and to predict 0 if the
112         pair was sampled from the marginals $p(\theta)p(x)$.
113         """
114 
115         assert theta.shape[0] == x.shape[0], "Batch sizes for theta and x must match."
116         batch_size = theta.shape[0]
117 
118         logits = self._classifier_logits(theta, x, num_atoms)
119         likelihood = torch.sigmoid(logits).squeeze()
120 
121         # Alternating pairs where there is one sampled from the joint and one
122         # sampled from the marginals. The first element is sampled from the
123         # joint p(theta, x) and is labelled 1. The second element is sampled
124         # from the marginals p(theta)p(x) and is labelled 0. And so on.
125         labels = ones(2 * batch_size)  # two atoms
126         labels[1::2] = 0.0
127 
128         # Binary cross entropy to learn the likelihood (AALR-specific)
129         return nn.BCELoss()(likelihood, labels)
130
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/sbi/inference/snre/snre_a.py b/sbi/inference/snre/snre_a.py
--- a/sbi/inference/snre/snre_a.py
+++ b/sbi/inference/snre/snre_a.py
@@ -122,7 +122,7 @@
         # sampled from the marginals. The first element is sampled from the
         # joint p(theta, x) and is labelled 1. The second element is sampled
         # from the marginals p(theta)p(x) and is labelled 0. And so on.
-        labels = ones(2 * batch_size)  # two atoms
+        labels = ones(2 * batch_size, device=self._device)  # two atoms
         labels[1::2] = 0.0
 
         # Binary cross entropy to learn the likelihood (AALR-specific)

{"golden_diff": "diff --git a/sbi/inference/snre/snre_a.py b/sbi/inference/snre/snre_a.py\n--- a/sbi/inference/snre/snre_a.py\n+++ b/sbi/inference/snre/snre_a.py\n@@ -122,7 +122,7 @@\n         # sampled from the marginals. The first element is sampled from the\n         # joint p(theta, x) and is labelled 1. The second element is sampled\n         # from the marginals p(theta)p(x) and is labelled 0. And so on.\n-        labels = ones(2 * batch_size)  # two atoms\n+        labels = ones(2 * batch_size, device=self._device)  # two atoms\n         labels[1::2] = 0.0\n \n         # Binary cross entropy to learn the likelihood (AALR-specific)\n", "issue": "Training SNRE_A on GPU fails\nHi! Using SNRE_A with `device=\"gpu\"` currently fails :-(\r\n\r\nThe error is as follows:\r\n```\r\nRuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!\r\n```\r\nThe origin of the issue is in `SNRE_A._loss` which instantiates `labels` without moving it to the device. Adding \r\n```\r\nlabels = labels.to(self._device)\r\n```\r\nbelow line 126 of `sbi/inference/snre/snre_a.py` fixes the issue.\n", "before_files": [{"content": "from typing import Any, Callable, Dict, Optional, Union\n\nimport torch\nfrom torch import Tensor, nn, ones\n\nfrom sbi.inference.posteriors.base_posterior import NeuralPosterior\nfrom sbi.inference.snre.snre_base import RatioEstimator\nfrom sbi.types import TensorboardSummaryWriter\nfrom sbi.utils import del_entries\n\n\nclass SNRE_A(RatioEstimator):\n    def __init__(\n        self,\n        prior,\n        classifier: Union[str, Callable] = \"resnet\",\n        device: str = \"cpu\",\n        logging_level: Union[int, str] = \"warning\",\n        summary_writer: Optional[TensorboardSummaryWriter] = None,\n        show_progress_bars: bool = True,\n        **unused_args\n    ):\n        r\"\"\"AALR[1], here known as SNRE_A.\n\n        [1] _Likelihood-free MCMC with Amortized Approximate Likelihood Ratios_, Hermans\n            et al., ICML 2020, https://arxiv.org/abs/1903.04057\n\n        Args:\n            prior: A probability distribution that expresses prior knowledge about the\n                parameters, e.g. which ranges are meaningful for them. Any\n                object with `.log_prob()`and `.sample()` (for example, a PyTorch\n                distribution) can be used.\n            classifier: Classifier trained to approximate likelihood ratios. If it is\n                a string, use a pre-configured network of the provided type (one of\n                linear, mlp, resnet). Alternatively, a function that builds a custom\n                neural network can be provided. The function will be called with the\n                first batch of simulations (theta, x), which can thus be used for shape\n                inference and potentially for z-scoring. It needs to return a PyTorch\n                `nn.Module` implementing the classifier.\n            device: torch device on which to compute, e.g. gpu, cpu.\n            logging_level: Minimum severity of messages to log. One of the strings\n                INFO, WARNING, DEBUG, ERROR and CRITICAL.\n            summary_writer: A tensorboard `SummaryWriter` to control, among others, log\n                file location (default is `<current working directory>/logs`.)\n            show_progress_bars: Whether to show a progressbar during simulation and\n                sampling.\n            unused_args: Absorbs additional arguments. No entries will be used. If it\n                is not empty, we warn. In future versions, when the new interface of\n                0.14.0 is more mature, we will remove this argument.\n        \"\"\"\n\n        kwargs = del_entries(locals(), entries=(\"self\", \"__class__\", \"unused_args\"))\n        super().__init__(**kwargs, **unused_args)\n\n    def train(\n        self,\n        training_batch_size: int = 50,\n        learning_rate: float = 5e-4,\n        validation_fraction: float = 0.1,\n        stop_after_epochs: int = 20,\n        max_num_epochs: Optional[int] = None,\n        clip_max_norm: Optional[float] = 5.0,\n        exclude_invalid_x: bool = True,\n        resume_training: bool = False,\n        discard_prior_samples: bool = False,\n        retrain_from_scratch_each_round: bool = False,\n        show_train_summary: bool = False,\n    ) -> NeuralPosterior:\n        r\"\"\"\n        Return classifier that approximates the ratio $p(\\theta,x)/p(\\theta)p(x)$.\n\n        Args:\n            training_batch_size: Training batch size.\n            learning_rate: Learning rate for Adam optimizer.\n            validation_fraction: The fraction of data to use for validation.\n            stop_after_epochs: The number of epochs to wait for improvement on the\n                validation set before terminating training.\n            max_num_epochs: Maximum number of epochs to run. If reached, we stop\n                training even when the validation loss is still decreasing. If None, we\n                train until validation loss increases (see also `stop_after_epochs`).\n            clip_max_norm: Value at which to clip the total gradient norm in order to\n                prevent exploding gradients. Use None for no clipping.\n            exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=\u00b1\u221e`\n                during training. Expect errors, silent or explicit, when `False`.\n            resume_training: Can be used in case training time is limited, e.g. on a\n                cluster. If `True`, the split between train and validation set, the\n                optimizer, the number of epochs, and the best validation log-prob will\n                be restored from the last time `.train()` was called.\n            discard_prior_samples: Whether to discard samples simulated in round 1, i.e.\n                from the prior. Training may be sped up by ignoring such less targeted\n                samples.\n            retrain_from_scratch_each_round: Whether to retrain the conditional density\n                estimator for the posterior from scratch each round.\n            show_train_summary: Whether to print the number of epochs and validation\n                loss and leakage after the training.\n\n        Returns:\n            Classifier that approximates the ratio $p(\\theta,x)/p(\\theta)p(x)$.\n        \"\"\"\n\n        # AALR is defined for `num_atoms=2`.\n        # Proxy to `super().__call__` to ensure right parameter.\n        kwargs = del_entries(locals(), entries=(\"self\", \"__class__\"))\n        return super().train(**kwargs, num_atoms=2)\n\n    def _loss(self, theta: Tensor, x: Tensor, num_atoms: int) -> Tensor:\n        \"\"\"\n        Returns the binary cross-entropy loss for the trained classifier.\n\n        The classifier takes as input a $(\\theta,x)$ pair. It is trained to predict 1\n        if the pair was sampled from the joint $p(\\theta,x)$, and to predict 0 if the\n        pair was sampled from the marginals $p(\\theta)p(x)$.\n        \"\"\"\n\n        assert theta.shape[0] == x.shape[0], \"Batch sizes for theta and x must match.\"\n        batch_size = theta.shape[0]\n\n        logits = self._classifier_logits(theta, x, num_atoms)\n        likelihood = torch.sigmoid(logits).squeeze()\n\n        # Alternating pairs where there is one sampled from the joint and one\n        # sampled from the marginals. The first element is sampled from the\n        # joint p(theta, x) and is labelled 1. The second element is sampled\n        # from the marginals p(theta)p(x) and is labelled 0. And so on.\n        labels = ones(2 * batch_size)  # two atoms\n        labels[1::2] = 0.0\n\n        # Binary cross entropy to learn the likelihood (AALR-specific)\n        return nn.BCELoss()(likelihood, labels)\n", "path": "sbi/inference/snre/snre_a.py"}], "after_files": [{"content": "from typing import Any, Callable, Dict, Optional, Union\n\nimport torch\nfrom torch import Tensor, nn, ones\n\nfrom sbi.inference.posteriors.base_posterior import NeuralPosterior\nfrom sbi.inference.snre.snre_base import RatioEstimator\nfrom sbi.types import TensorboardSummaryWriter\nfrom sbi.utils import del_entries\n\n\nclass SNRE_A(RatioEstimator):\n    def __init__(\n        self,\n        prior,\n        classifier: Union[str, Callable] = \"resnet\",\n        device: str = \"cpu\",\n        logging_level: Union[int, str] = \"warning\",\n        summary_writer: Optional[TensorboardSummaryWriter] = None,\n        show_progress_bars: bool = True,\n        **unused_args\n    ):\n        r\"\"\"AALR[1], here known as SNRE_A.\n\n        [1] _Likelihood-free MCMC with Amortized Approximate Likelihood Ratios_, Hermans\n            et al., ICML 2020, https://arxiv.org/abs/1903.04057\n\n        Args:\n            prior: A probability distribution that expresses prior knowledge about the\n                parameters, e.g. which ranges are meaningful for them. Any\n                object with `.log_prob()`and `.sample()` (for example, a PyTorch\n                distribution) can be used.\n            classifier: Classifier trained to approximate likelihood ratios. If it is\n                a string, use a pre-configured network of the provided type (one of\n                linear, mlp, resnet). Alternatively, a function that builds a custom\n                neural network can be provided. The function will be called with the\n                first batch of simulations (theta, x), which can thus be used for shape\n                inference and potentially for z-scoring. It needs to return a PyTorch\n                `nn.Module` implementing the classifier.\n            device: torch device on which to compute, e.g. gpu, cpu.\n            logging_level: Minimum severity of messages to log. One of the strings\n                INFO, WARNING, DEBUG, ERROR and CRITICAL.\n            summary_writer: A tensorboard `SummaryWriter` to control, among others, log\n                file location (default is `<current working directory>/logs`.)\n            show_progress_bars: Whether to show a progressbar during simulation and\n                sampling.\n            unused_args: Absorbs additional arguments. No entries will be used. If it\n                is not empty, we warn. In future versions, when the new interface of\n                0.14.0 is more mature, we will remove this argument.\n        \"\"\"\n\n        kwargs = del_entries(locals(), entries=(\"self\", \"__class__\", \"unused_args\"))\n        super().__init__(**kwargs, **unused_args)\n\n    def train(\n        self,\n        training_batch_size: int = 50,\n        learning_rate: float = 5e-4,\n        validation_fraction: float = 0.1,\n        stop_after_epochs: int = 20,\n        max_num_epochs: Optional[int] = None,\n        clip_max_norm: Optional[float] = 5.0,\n        exclude_invalid_x: bool = True,\n        resume_training: bool = False,\n        discard_prior_samples: bool = False,\n        retrain_from_scratch_each_round: bool = False,\n        show_train_summary: bool = False,\n    ) -> NeuralPosterior:\n        r\"\"\"\n        Return classifier that approximates the ratio $p(\\theta,x)/p(\\theta)p(x)$.\n\n        Args:\n            training_batch_size: Training batch size.\n            learning_rate: Learning rate for Adam optimizer.\n            validation_fraction: The fraction of data to use for validation.\n            stop_after_epochs: The number of epochs to wait for improvement on the\n                validation set before terminating training.\n            max_num_epochs: Maximum number of epochs to run. If reached, we stop\n                training even when the validation loss is still decreasing. If None, we\n                train until validation loss increases (see also `stop_after_epochs`).\n            clip_max_norm: Value at which to clip the total gradient norm in order to\n                prevent exploding gradients. Use None for no clipping.\n            exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=\u00b1\u221e`\n                during training. Expect errors, silent or explicit, when `False`.\n            resume_training: Can be used in case training time is limited, e.g. on a\n                cluster. If `True`, the split between train and validation set, the\n                optimizer, the number of epochs, and the best validation log-prob will\n                be restored from the last time `.train()` was called.\n            discard_prior_samples: Whether to discard samples simulated in round 1, i.e.\n                from the prior. Training may be sped up by ignoring such less targeted\n                samples.\n            retrain_from_scratch_each_round: Whether to retrain the conditional density\n                estimator for the posterior from scratch each round.\n            show_train_summary: Whether to print the number of epochs and validation\n                loss and leakage after the training.\n\n        Returns:\n            Classifier that approximates the ratio $p(\\theta,x)/p(\\theta)p(x)$.\n        \"\"\"\n\n        # AALR is defined for `num_atoms=2`.\n        # Proxy to `super().__call__` to ensure right parameter.\n        kwargs = del_entries(locals(), entries=(\"self\", \"__class__\"))\n        return super().train(**kwargs, num_atoms=2)\n\n    def _loss(self, theta: Tensor, x: Tensor, num_atoms: int) -> Tensor:\n        \"\"\"\n        Returns the binary cross-entropy loss for the trained classifier.\n\n        The classifier takes as input a $(\\theta,x)$ pair. It is trained to predict 1\n        if the pair was sampled from the joint $p(\\theta,x)$, and to predict 0 if the\n        pair was sampled from the marginals $p(\\theta)p(x)$.\n        \"\"\"\n\n        assert theta.shape[0] == x.shape[0], \"Batch sizes for theta and x must match.\"\n        batch_size = theta.shape[0]\n\n        logits = self._classifier_logits(theta, x, num_atoms)\n        likelihood = torch.sigmoid(logits).squeeze()\n\n        # Alternating pairs where there is one sampled from the joint and one\n        # sampled from the marginals. The first element is sampled from the\n        # joint p(theta, x) and is labelled 1. The second element is sampled\n        # from the marginals p(theta)p(x) and is labelled 0. And so on.\n        labels = ones(2 * batch_size, device=self._device)  # two atoms\n        labels[1::2] = 0.0\n\n        # Binary cross entropy to learn the likelihood (AALR-specific)\n        return nn.BCELoss()(likelihood, labels)\n", "path": "sbi/inference/snre/snre_a.py"}]}