As @albanD suggested try considering SGD instead of Adam as Adam needs more memory for storing first and second moment vectors.
Also, reducing embedding size and bucket_size and overall architecture might help.
Also want to suggest you that, bucketize() only needs tensor_shape instead of the whole tensor.