Multi-GPU TFF simulation - gpu

Question

When running my custom federated learning code on GPU (using tff.templates.IterativeProcess): I am getting following warning:
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2022-11-12 12:22:28.981770: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1532] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 208 MB memory: -> device: 0, name: NVIDIA GeForce RTX 3080, pci bus id: 0000:1b:00.0, compute capability: 8.6
2022-11-12 12:22:28.982594: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1532] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 7701 MB memory: -> device: 1, name: NVIDIA GeForce RTX 3080, pci bus id: 0000:1e:00.0, compute capability: 8.6
2022-11-12 12:22:28.983096: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1532] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 7701 MB memory: -> device: 2, name: NVIDIA GeForce RTX 3080, pci bus id: 0000:3f:00.0, compute capability: 8.6
2022-11-12 12:22:28.983653: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1532] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 7701 MB memory: -> device: 3, name: NVIDIA GeForce RTX 3080, pci bus id: 0000:40:00.0, compute capability: 8.6
and then following error while calling tff.templates.IterativeProcess for training:
Traceback (most recent call last):
File "/home/user1/p1/cases/a_ii.py", line 263, in <module>
server_parameters = FL.next(server_parameters, rate, federated_train_data, tff_val(sel))
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/computation/computation_impl.py", line 130, in __call__
return self._context_stack.current.invoke(self, arg)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/execution_contexts/sync_execution_context.py", line 67, in invoke
return self._async_runner.run_coro_and_return_result(
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/async_utils.py", line 223, in run_coro_and_return_result
return future.result()
File "/home/user1/anaconda3/lib/python3.9/concurrent/futures/_base.py", line 446, in result
return self.__get_result()
File "/home/user1/anaconda3/lib/python3.9/concurrent/futures/_base.py", line 391, in __get_result
raise self._exception
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 258, in __step
result = coro.throw(exc)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/retrying.py", line 107, in retry_coro_fn
raise e
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/retrying.py", line 97, in retry_coro_fn
result = await fn(*args, **kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/execution_contexts/async_execution_context.py", line 231, in invoke
return await tracing.wrap_coroutine_in_current_trace_context(
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 391, in _wrapped
return await coro
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/execution_contexts/async_execution_context.py", line 141, in _invoke
result = await executor.create_call(comp, arg)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 346, in create_call
return await comp_repr.invoke(self, arg)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 166, in invoke
return await executor._evaluate(comp_lambda.result, new_scope) # pylint: disable=protected-access
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 516, in _evaluate
return await self._evaluate_block(comp, scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 480, in _evaluate_block
return await self._evaluate(comp.block.result, scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 508, in _evaluate
return await self._evaluate_reference(comp, scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 432, in _evaluate_reference
return await scope.resolve_reference(comp.reference.name)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 115, in resolve_reference
return await value
File "/home/user1/anaconda3/lib/python3.9/asyncio/futures.py", line 284, in __await__
yield self # This tells Task to wait for completion.
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 328, in __wakeup
future.result()
File "/home/user1/anaconda3/lib/python3.9/asyncio/futures.py", line 201, in result
raise self._exception
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 258, in __step
result = coro.throw(exc)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 510, in _evaluate
return await self._evaluate_call(comp, scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 448, in _evaluate_call
func, arg = await asyncio.gather(func, get_arg())
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 328, in __wakeup
future.result()
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 258, in __step
result = coro.throw(exc)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 444, in get_arg
return await self._evaluate(comp.call.argument, scope=scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 514, in _evaluate
return await self._evaluate_struct(comp, scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 468, in _evaluate_struct
values = await asyncio.gather(*values)
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 328, in __wakeup
future.result()
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 258, in __step
result = coro.throw(exc)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 508, in _evaluate
return await self._evaluate_reference(comp, scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 432, in _evaluate_reference
return await scope.resolve_reference(comp.reference.name)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 115, in resolve_reference
return await value
File "/home/user1/anaconda3/lib/python3.9/asyncio/futures.py", line 284, in __await__
yield self # This tells Task to wait for completion.
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 328, in __wakeup
future.result()
File "/home/user1/anaconda3/lib/python3.9/asyncio/futures.py", line 201, in result
raise self._exception
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 258, in __step
result = coro.throw(exc)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 508, in _evaluate
return await self._evaluate_reference(comp, scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 432, in _evaluate_reference
return await scope.resolve_reference(comp.reference.name)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 117, in resolve_reference
return await self._parent.resolve_reference(name)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 115, in resolve_reference
return await value
File "/home/user1/anaconda3/lib/python3.9/asyncio/futures.py", line 284, in __await__
yield self # This tells Task to wait for completion.
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 328, in __wakeup
future.result()
File "/home/user1/anaconda3/lib/python3.9/asyncio/futures.py", line 201, in result
raise self._exception
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 258, in __step
result = coro.throw(exc)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 510, in _evaluate
return await self._evaluate_call(comp, scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 448, in _evaluate_call
func, arg = await asyncio.gather(func, get_arg())
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 328, in __wakeup
future.result()
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 258, in __step
result = coro.throw(exc)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 510, in _evaluate
return await self._evaluate_call(comp, scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 448, in _evaluate_call
func, arg = await asyncio.gather(func, get_arg())
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 328, in __wakeup
future.result()
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 258, in __step
result = coro.throw(exc)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 444, in get_arg
return await self._evaluate(comp.call.argument, scope=scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 514, in _evaluate
return await self._evaluate_struct(comp, scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 468, in _evaluate_struct
values = await asyncio.gather(*values)
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 328, in __wakeup
future.result()
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 258, in __step
result = coro.throw(exc)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 444, in get_arg
return await self._evaluate(comp.call.argument, scope=scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 514, in _evaluate
return await self._evaluate_struct(comp, scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 468, in _evaluate_struct
values = await asyncio.gather(*values)
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 328, in __wakeup
future.result()
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 258, in __step
result = coro.throw(exc)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 508, in _evaluate
return await self._evaluate_reference(comp, scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 432, in _evaluate_reference
return await scope.resolve_reference(comp.reference.name)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 117, in resolve_reference
return await self._parent.resolve_reference(name)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 115, in resolve_reference
return await value
File "/home/user1/anaconda3/lib/python3.9/asyncio/futures.py", line 284, in __await__
yield self # This tells Task to wait for completion.
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 328, in __wakeup
future.result()
File "/home/user1/anaconda3/lib/python3.9/asyncio/futures.py", line 201, in result
raise self._exception
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 258, in __step
result = coro.throw(exc)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 508, in _evaluate
return await self._evaluate_reference(comp, scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 432, in _evaluate_reference
return await scope.resolve_reference(comp.reference.name)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 115, in resolve_reference
return await value
File "/home/user1/anaconda3/lib/python3.9/asyncio/futures.py", line 284, in __await__
yield self # This tells Task to wait for completion.
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 328, in __wakeup
future.result()
File "/home/user1/anaconda3/lib/python3.9/asyncio/futures.py", line 201, in result
raise self._exception
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 258, in __step
result = coro.throw(exc)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 510, in _evaluate
return await self._evaluate_call(comp, scope)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 449, in _evaluate_call
return await self.create_call(func, arg=arg)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/reference_resolving_executor.py", line 342, in create_call
return ReferenceResolvingExecutorValue(await
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/thread_delegating_executor.py", line 102, in create_call
return await self._delegate(self._target_executor.create_call(comp, arg))
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/thread_delegating_executor.py", line 87, in _delegate
result_value = await _delegate_with_trace_ctx(coro, self._async_runner)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/async_utils.py", line 227, in await_coro_and_return_result
return await asyncio.wrap_future(
File "/home/user1/anaconda3/lib/python3.9/asyncio/futures.py", line 284, in __await__
yield self # This tells Task to wait for completion.
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 328, in __wakeup
future.result()
File "/home/user1/anaconda3/lib/python3.9/asyncio/futures.py", line 201, in result
raise self._exception
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 258, in __step
result = coro.throw(exc)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 391, in _wrapped
return await coro
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/federating_executor.py", line 457, in create_call
return await self._strategy.compute_federated_intrinsic(
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/federating_executor.py", line 143, in compute_federated_intrinsic
return await fn(arg) # pylint: disable=not-callable
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/federated_resolving_strategy.py", line 459, in compute_federated_map
return await self._map(arg, all_equal=False)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/federated_resolving_strategy.py", line 340, in _map
results = await asyncio.gather(*[
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 328, in __wakeup
future.result()
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 258, in __step
result = coro.throw(exc)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/federated_resolving_strategy.py", line 337, in _map_child
fn_at_child = await child.create_value(fn, fn_type)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/thread_delegating_executor.py", line 92, in create_value
return await self._delegate(
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/thread_delegating_executor.py", line 87, in _delegate
result_value = await _delegate_with_trace_ctx(coro, self._async_runner)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/async_utils.py", line 227, in await_coro_and_return_result
return await asyncio.wrap_future(
File "/home/user1/anaconda3/lib/python3.9/asyncio/futures.py", line 284, in __await__
yield self # This tells Task to wait for completion.
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 328, in __wakeup
future.result()
File "/home/user1/anaconda3/lib/python3.9/asyncio/futures.py", line 201, in result
raise self._exception
File "/home/user1/anaconda3/lib/python3.9/asyncio/tasks.py", line 256, in __step
result = coro.send(None)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 391, in _wrapped
return await coro
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 201, in async_trace
result = await fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/eager_tf_executor.py", line 682, in create_value
normalized_value = to_representation_for_type(value,
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 228, in sync_trace
result = fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/eager_tf_executor.py", line 518, in to_representation_for_type
return _to_computation_internal_rep(
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 228, in sync_trace
result = fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/eager_tf_executor.py", line 404, in _to_computation_internal_rep
embedded_fn = embed_tensorflow_computation(value, type_spec, device)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/common_libs/tracing.py", line 228, in sync_trace
result = fn(*fn_args, **fn_kwargs)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/eager_tf_executor.py", line 272, in embed_tensorflow_computation
comp = _ensure_comp_runtime_compatible(comp)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/eager_tf_executor.py", line 245, in _ensure_comp_runtime_compatible
_check_dataset_reduce_for_multi_gpu(graph_def)
File "/home/user1/fl/lib/python3.9/site-packages/tensorflow_federated/python/core/impl/executors/eager_tf_executor.py", line 63, in _check_dataset_reduce_for_multi_gpu
raise ValueError(
ValueError: Detected dataset reduce op in multi-GPU TFF simulation: `use_experimental_simulation_loop=True` for `tff.learning`; or use `for ... in iter(dataset)` for your own dataset iterations. See https://www.tensorflow.org/federated/tutorials/simulations_with_accelerators for examples.
Earlier, the code was working fine, but after GPU rebooted, code started giving warnings and errors. Using use_experimental_simulation_loop=True in tff.learning works fine, but in tff.templates.IterativeProcess, there is no such argument. Does anyone know how to solve this problem?