





'''simple docstring'''






from collections import deque







class 			lowerCamelCase					:


					def __init__(					self	:  Any       ,     __snake_case	:  str       ,     __snake_case	:  int       ,     __snake_case	:  int       ):





											'''simple docstring'''





											_snake_case:   Any  					=		process_name  # process name
											_snake_case:   Optional[Any]  					=		arrival_time  # arrival time of the process
											# completion time of finished process or last interrupted time
											_snake_case:   Union[str, Any]  					=		arrival_time
											_snake_case:   Any  					=		burst_time  # remaining burst time
											_snake_case:   List[str]  					=		0  # total time of the process wait in ready queue
											_snake_case:   Union[str, Any]  					=		0  # time from arrival time to completion time







class 			lowerCamelCase					:


					def __init__(					self	:  str       ,     __snake_case	:  int       ,     __snake_case	:  list[int]       ,     __snake_case	:  deque[Process]       ,     __snake_case	:  int       ,     ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		number_of_queues
											# time slice of queues that round robin algorithm applied
											_snake_case:   Optional[int]  					=		time_slices
											# unfinished process is in this ready_queue
											_snake_case:   Tuple  					=		queue
											# current time
											_snake_case:   Tuple  					=		current_time
											# finished process is in this sequence queue
											_snake_case:   deque[Process]  					=		deque()


					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ):





											'''simple docstring'''





											_snake_case:   int  					=		[]
											for i in range(len(self.finish_queue       )       ):
																	sequence.append(self.finish_queue[i].process_name       )
											return sequence


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[Any]       ,     __snake_case	:  list[Process]       ):





											'''simple docstring'''





											_snake_case:   Dict  					=		[]
											for i in range(len(__snake_case       )       ):
																	waiting_times.append(queue[i].waiting_time       )
											return waiting_times


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ,     __snake_case	:  list[Process]       ):





											'''simple docstring'''





											_snake_case:   Union[str, Any]  					=		[]
											for i in range(len(__snake_case       )       ):
																	turnaround_times.append(queue[i].turnaround_time       )
											return turnaround_times


					def    SCREAMING_SNAKE_CASE_		(					self	:  str       ,     __snake_case	:  list[Process]       ):





											'''simple docstring'''





											_snake_case:   Optional[int]  					=		[]
											for i in range(len(__snake_case       )       ):
																	completion_times.append(queue[i].stop_time       )
											return completion_times


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[Any]       ,     __snake_case	:  deque[Process]       ):





											'''simple docstring'''





											return [q.burst_time for q in queue]


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[int]       ,     __snake_case	:  Process       ):





											'''simple docstring'''





											process.waiting_time += self.current_time - process.stop_time
											return process.waiting_time


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ,     __snake_case	:  deque[Process]       ):





											'''simple docstring'''





											_snake_case:   deque[Process]  					=		deque()  # sequence deque of finished process
											while len(__snake_case       ) != 0:
																	_snake_case:   Any  					=		ready_queue.popleft()  # current process

																	# if process's arrival time is later than current time, update current time
																	if self.current_time < cp.arrival_time:
																							self.current_time += cp.arrival_time

																	# update waiting time of current process
																	self.update_waiting_time(__snake_case       )
																	# update current time
																	self.current_time += cp.burst_time
																	# finish the process and set the process's burst-time 0
																	_snake_case:   int  					=		0
																	# set the process's turnaround time because it is finished
																	_snake_case:   Any  					=		self.current_time - cp.arrival_time
																	# set the completion time
																	_snake_case:   int  					=		self.current_time
																	# add the process to queue that has finished queue
																	finished.append(__snake_case       )

											self.finish_queue.extend(__snake_case       )  # add finished process to finish queue
											# FCFS will finish all remaining processes
											return finished


					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ,     __snake_case	:  deque[Process]       ,     __snake_case	:  int       ):





											'''simple docstring'''





											_snake_case:   deque[Process]  					=		deque()  # sequence deque of terminated process
											# just for 1 cycle and unfinished processes will go back to queue
											for _ in range(len(__snake_case       )       ):
																	_snake_case:   Tuple  					=		ready_queue.popleft()  # current process

																	# if process's arrival time is later than current time, update current time
																	if self.current_time < cp.arrival_time:
																							self.current_time += cp.arrival_time

																	# update waiting time of unfinished processes
																	self.update_waiting_time(__snake_case       )
																	# if the burst time of process is bigger than time-slice
																	if cp.burst_time > time_slice:
																							# use CPU for only time-slice
																							self.current_time += time_slice
																							# update remaining burst time
																							cp.burst_time -= time_slice
																							# update end point time
																							_snake_case:   Optional[int]  					=		self.current_time
																							# locate the process behind the queue because it is not finished
																							ready_queue.append(__snake_case       )
																	else:
																							# use CPU for remaining burst time
																							self.current_time += cp.burst_time
																							# set burst time 0 because the process is finished
																							_snake_case:   Dict  					=		0
																							# set the finish time
																							_snake_case:   Tuple  					=		self.current_time
																							# update the process' turnaround time because it is finished
																							_snake_case:   Any  					=		self.current_time - cp.arrival_time
																							# add the process to queue that has finished queue
																							finished.append(__snake_case       )

											self.finish_queue.extend(__snake_case       )  # add finished process to finish queue
											# return finished processes queue and remaining processes queue
											return finished, ready_queue




					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											for i in range(self.number_of_queues - 1       ):
																	_snake_case					,						_snake_case:   List[str]  					=		self.round_robin(
																	    self.ready_queue       ,     self.time_slices[i]       )
											#  the last queue has first_come_first_served algorithm
											self.first_come_first_served(self.ready_queue       )

											return self.finish_queue


if __name__ == "__main__":
		import doctest

		A   : Tuple             =			Process('P1', 0, 53)
		A   : List[Any]             =			Process('P2', 0, 17)
		A   : Tuple             =			Process('P3', 0, 68)
		A   : str             =			Process('P4', 0, 24)
		A   : str             =			3
		A   : Tuple             =			[17, 25]
		A   : Optional[Any]             =			deque([Pa, Pa, Pa, Pa])

		if len(time_slices) != number_of_queues - 1:
				raise SystemExit(0)

		doctest.testmod(extraglobs={'queue': deque([Pa, Pa, Pa, Pa])})

		A   : Union[str, Any]             =			Process('P1', 0, 53)
		A   : Union[str, Any]             =			Process('P2', 0, 17)
		A   : Tuple             =			Process('P3', 0, 68)
		A   : Optional[int]             =			Process('P4', 0, 24)
		A   : Optional[int]             =			3
		A   : int             =			[17, 25]
		A   : List[str]             =			deque([Pa, Pa, Pa, Pa])
		A   : Tuple             =			MLFQ(number_of_queues, time_slices, queue, 0)
		A   : Optional[Any]             =			mlfq.multi_level_feedback_queue()

		# print total waiting times of processes(P1, P2, P3, P4)
		print(
		    F'waiting time:\
        \t\t\t{MLFQ.calculate_waiting_time(mlfq, [Pa, Pa, Pa, Pa])}'
		)
		# print completion times of processes(P1, P2, P3, P4)
		print(
		    F'completion time:\
        \t\t{MLFQ.calculate_completion_time(mlfq, [Pa, Pa, Pa, Pa])}'
		)
		# print total turnaround times of processes(P1, P2, P3, P4)
		print(
		    F'turnaround time:\
        \t\t{MLFQ.calculate_turnaround_time(mlfq, [Pa, Pa, Pa, Pa])}'
		)
		# print sequence of finished processes
		print(
		    F'sequence of finished processes:\
        {mlfq.calculate_sequence_of_finish_queue()}'
		)












'''simple docstring'''






import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu







class 			lowerCamelCase					(						nn.Module  ):


					def __init__(					self	:  Optional[int]       ):





											'''simple docstring'''





											super().__init__()
											_snake_case:   int  					=		nn.Linear(3       ,     4       )
											_snake_case:   Any  					=		nn.BatchNormad(4       )
											_snake_case:   Union[str, Any]  					=		nn.Linear(4       ,     5       )




					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ,     __snake_case	:  List[Any]       ):





											'''simple docstring'''





											return self.lineara(self.batchnorm(self.lineara(__snake_case       )       )       )







class 			lowerCamelCase					(						__UpperCAmelCase  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ,     __snake_case	:  Optional[Any]       ,     *__snake_case	:  Optional[Any]       ,     **__snake_case	:  List[Any]       ):





											'''simple docstring'''





											return (args[0] + 1,) + args[1:], kwargs







class 			lowerCamelCase					(						__UpperCAmelCase  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  Tuple       ):





											'''simple docstring'''





											return output + 1







class 			lowerCamelCase					(						unittest.TestCase  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ):





											'''simple docstring'''





											_snake_case:   str  					=		ModelForTest()
											_snake_case:   List[str]  					=		ModelHook()

											add_hook_to_module(__snake_case       ,     __snake_case       )
											self.assertEqual(test_model._hf_hook       ,     __snake_case       )
											self.assertTrue(hasattr(__snake_case       ,     '_old_forward'       )       )

											# Check adding the hook did not change the name or the signature
											self.assertEqual(test_model.forward.__name__       ,     'forward'       )
											self.assertListEqual(list(inspect.signature(test_model.forward       ).parameters       )       ,     ['x']       )

											remove_hook_from_module(__snake_case       )
											self.assertFalse(hasattr(__snake_case       ,     '_hf_hook'       )       )
											self.assertFalse(hasattr(__snake_case       ,     '_old_forward'       )       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[Any]       ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		ModelForTest()
											_snake_case:   Union[str, Any]  					=		ModelHook()

											add_hook_to_module(__snake_case       ,     __snake_case       )
											add_hook_to_module(__snake_case       ,     __snake_case       ,     append=__snake_case       )

											self.assertEqual(isinstance(test_model._hf_hook       ,     __snake_case       )       ,     __snake_case       )
											self.assertEqual(len(test_model._hf_hook.hooks       )       ,     2       )
											self.assertTrue(hasattr(__snake_case       ,     '_old_forward'       )       )

											# Check adding the hook did not change the name or the signature
											self.assertEqual(test_model.forward.__name__       ,     'forward'       )
											self.assertListEqual(list(inspect.signature(test_model.forward       ).parameters       )       ,     ['x']       )

											remove_hook_from_module(__snake_case       )
											self.assertFalse(hasattr(__snake_case       ,     '_hf_hook'       )       )
											self.assertFalse(hasattr(__snake_case       ,     '_old_forward'       )       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





											_snake_case:   Dict  					=		ModelForTest()
											_snake_case:   Tuple  					=		torch.randn(2       ,     3       )
											_snake_case:   List[Any]  					=		test_model(x + 1       )
											_snake_case:   int  					=		test_model(x + 2       )

											_snake_case:   List[str]  					=		PreForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   List[Any]  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     __snake_case       ,     atol=1e-5       )       )

											# Attaching a hook to a model when it already has one replaces, does not chain
											_snake_case:   Any  					=		PreForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   Dict  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     __snake_case       ,     atol=1e-5       )       )

											# You need to use the sequential hook to chain two or more hooks
											_snake_case:   Any  					=		SequentialHook(PreForwardHook()       ,     PreForwardHook()       )
											add_hook_to_module(__snake_case       ,     __snake_case       )

											_snake_case:   int  					=		test_model(__snake_case       )
											assert torch.allclose(__snake_case       ,     __snake_case       ,     atol=1e-5       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											_snake_case:   int  					=		ModelForTest()
											_snake_case:   str  					=		torch.randn(2       ,     3       )
											_snake_case:   int  					=		test_model(__snake_case       )

											_snake_case:   Union[str, Any]  					=		PostForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   int  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     output + 1       ,     atol=1e-5       )       )

											# Attaching a hook to a model when it already has one replaces, does not chain
											_snake_case:   Tuple  					=		PostForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   Optional[int]  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     output + 1       ,     atol=1e-5       )       )

											# You need to use the sequential hook to chain two or more hooks
											_snake_case:   List[str]  					=		SequentialHook(PostForwardHook()       ,     PostForwardHook()       )
											add_hook_to_module(__snake_case       ,     __snake_case       )

											_snake_case:   List[Any]  					=		test_model(__snake_case       )
											assert torch.allclose(__snake_case       ,     output + 2       ,     atol=1e-5       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		ModelForTest()
											_snake_case:   Dict  					=		torch.randn(2       ,     3       )
											_snake_case:   Tuple  					=		test_model(__snake_case       )

											_snake_case:   Tuple  					=		PostForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   List[str]  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     output + 1       )       )
											self.assertTrue(outputa.requires_grad       )

											_snake_case:   Any  					=		True
											_snake_case:   str  					=		test_model(__snake_case       )
											self.assertFalse(outputa.requires_grad       )


					@require_multi_gpu
					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											_snake_case:   Any  					=		ModelForTest()
											# Everything is on CPU
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )

											# This will move each submodule on different devices
											add_hook_to_module(model.lineara       ,     AlignDevicesHook(execution_device=0       )       )
											add_hook_to_module(model.batchnorm       ,     AlignDevicesHook(execution_device=0       )       )
											add_hook_to_module(model.lineara       ,     AlignDevicesHook(execution_device=1       )       )

											self.assertEqual(model.lineara.weight.device       ,     torch.device(0       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device(0       )       )
											self.assertEqual(model.batchnorm.running_mean.device       ,     torch.device(0       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device(1       )       )

											# We can still make a forward pass. The input does not need to be on any particular device
											_snake_case:   int  					=		torch.randn(2       ,     3       )
											_snake_case:   Tuple  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     torch.device(1       )       )

											# We can add a general hook to put back output on same device as input.
											add_hook_to_module(__snake_case       ,     AlignDevicesHook(io_same_device=__snake_case       )       )
											_snake_case:   Optional[Any]  					=		torch.randn(2       ,     3       ).to(0       )
											_snake_case:   Any  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     torch.device(0       )       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case:   Union[str, Any]  					=		ModelForTest()

											# Everything is on CPU
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )

											# This will move each submodule on different devices
											_snake_case:   Optional[int]  					=		{'execution_device': 0 if torch.cuda.is_available() else 'cpu', 'offload': True}

											add_hook_to_module(model.lineara       ,     AlignDevicesHook(**__snake_case       )       )
											add_hook_to_module(model.batchnorm       ,     AlignDevicesHook(**__snake_case       )       )
											add_hook_to_module(model.lineara       ,     AlignDevicesHook(**__snake_case       )       )

											# Parameters have been offloaded, so on the meta device
											self.assertEqual(model.lineara.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('meta'       )       )
											# Buffers are not included in the offload by default, so are on the execution device
											_snake_case:   Optional[int]  					=		torch.device(hook_kwargs['execution_device']       )
											self.assertEqual(model.batchnorm.running_mean.device       ,     __snake_case       )

											_snake_case:   Tuple  					=		torch.randn(2       ,     3       )
											_snake_case:   Union[str, Any]  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_module(model.lineara       )
											remove_hook_from_module(model.batchnorm       )
											remove_hook_from_module(model.lineara       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )

											# Now test with buffers included in the offload
											_snake_case:   Optional[Any]  					=		{
											    'execution_device': 0 if torch.cuda.is_available() else 'cpu',
											    'offload': True,
											    'offload_buffers': True,
											}

											add_hook_to_module(model.lineara       ,     AlignDevicesHook(**__snake_case       )       )
											add_hook_to_module(model.batchnorm       ,     AlignDevicesHook(**__snake_case       )       )
											add_hook_to_module(model.lineara       ,     AlignDevicesHook(**__snake_case       )       )

											# Parameters have been offloaded, so on the meta device, buffers included
											self.assertEqual(model.lineara.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.running_mean.device       ,     torch.device('meta'       )       )

											_snake_case:   Optional[Any]  					=		torch.randn(2       ,     3       )
											_snake_case:   Union[str, Any]  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_module(model.lineara       )
											remove_hook_from_module(model.batchnorm       )
											remove_hook_from_module(model.lineara       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case:   Union[str, Any]  					=		ModelForTest()

											# Everything is on CPU
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )

											# This will move each submodule on different devices
											_snake_case:   Optional[Any]  					=		0 if torch.cuda.is_available() else 'cpu'
											attach_align_device_hook(__snake_case       ,     execution_device=__snake_case       ,     offload=__snake_case       )

											# Parameters have been offloaded, so on the meta device
											self.assertEqual(model.lineara.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('meta'       )       )
											# Buffers are not included in the offload by default, so are on the execution device
											_snake_case:   List[str]  					=		torch.device(__snake_case       )
											self.assertEqual(model.batchnorm.running_mean.device       ,     __snake_case       )

											_snake_case:   Optional[Any]  					=		torch.randn(2       ,     3       )
											_snake_case:   List[str]  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_submodules(__snake_case       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )

											# Now test with buffers included in the offload
											attach_align_device_hook(__snake_case       ,     execution_device=__snake_case       ,     offload=__snake_case       ,     offload_buffers=__snake_case       )

											# Parameters have been offloaded, so on the meta device, buffers included
											self.assertEqual(model.lineara.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.running_mean.device       ,     torch.device('meta'       )       )

											_snake_case:   Any  					=		torch.randn(2       ,     3       )
											_snake_case:   List[Any]  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_submodules(__snake_case       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )




					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case:   Tuple  					=		ModelForTest()

											# Everything is on CPU
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )

											# This will move each submodule on different devices
											_snake_case:   Tuple  					=		0 if torch.cuda.is_available() else 'cpu'
											attach_align_device_hook(
											    __snake_case       ,     execution_device=__snake_case       ,     offload=__snake_case       ,     weights_map=model.state_dict()       )

											# Parameters have been offloaded, so on the meta device
											self.assertEqual(model.lineara.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('meta'       )       )
											# Buffers are not included in the offload by default, so are on the execution device
											_snake_case:   str  					=		torch.device(__snake_case       )
											self.assertEqual(model.batchnorm.running_mean.device       ,     __snake_case       )

											_snake_case:   Dict  					=		torch.randn(2       ,     3       )
											_snake_case:   Dict  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_submodules(__snake_case       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )

											# Now test with buffers included in the offload
											attach_align_device_hook(
											    __snake_case       ,     execution_device=__snake_case       ,     offload=__snake_case       ,     weights_map=model.state_dict()       ,     offload_buffers=__snake_case       ,     )

											# Parameters have been offloaded, so on the meta device, buffers included
											self.assertEqual(model.lineara.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('meta'       )       )
											self.assertEqual(model.batchnorm.running_mean.device       ,     torch.device('meta'       )       )

											_snake_case:   List[Any]  					=		torch.randn(2       ,     3       )
											_snake_case:   List[str]  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_submodules(__snake_case       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.batchnorm.weight.device       ,     torch.device('cpu'       )       )
											self.assertEqual(model.lineara.weight.device       ,     torch.device('cpu'       )       )





