localize tensorflow import in construct_batch

PennyLaneAI · May 31, 2024 · 14696c8 · 14696c8
1 parent c7abcf2
commit 14696c8
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/pennylane/workflow/construct_batch.py b/pennylane/workflow/construct_batch.py
@@ -18,8 +18,6 @@
 from functools import wraps
 from typing import Callable, Literal, Optional, Tuple, Union
 
-import tensorflow as tf
-
 import pennylane as qml
 
 from .qnode import QNode, _get_device_shots, _make_execution_config
@@ -299,6 +297,9 @@ def batch_constructor(*args, **kwargs) -> Tuple[Tuple["qml.tape.QuantumTape", Ca
  shots = kwargs.pop("shots", _get_device_shots(qnode.device))
 
  if isinstance(qnode, qml.qnn.KerasLayer):
+ # pylint: disable=import-outside-toplevel
+ import tensorflow as tf
+
  with tf.GradientTape() as tape:
  tape.watch(list(qnode.qnode_weights.values()))